Skill: bahuia/heteromind (version 0.3.0)

--- name: heteromind description: Unified heterogeneous knowledge QA system. Automatically routes natural language queries to SQL databases, Knowledge Graphs, or table files using 4-layer detection (rule-based, LLM semantic, schema matching, entity verification). Supports multi-LLM providers and bilingual queries. Trigger on data queries, "how many", "show", aggregations, filters, joins, or structured information requests. required_env_vars: - DEEPSEEK_API_KEY - OPENAI_API_KEY optional_env_vars: - MYSQL_CONNECTION_STRING - POSTGRES_CONNECTION_STRING - CUSTOM_KG_ENDPOINT - WORKSPACE - TABLE_PATHS --- # HeteroMind Unified heterogeneous knowledge QA system with automatic source detection and multi-stage reasoning. ## Core Concept Natural language queries are automatically routed to the appropriate knowledge source (SQL, Knowledge Graph, or Table files) without requiring users to specify the data source. A 4-layer detection architecture ensures accurate source identification, followed by multi-stage query generation with self-revision and voting. ``` User Query → Source Detection (4 layers) → Query Generation → Self-Revision → Voting → Execution → Answer ``` ## When to Use | Trigger | Action | |---------|--------| | "How many employees in X?" | NL2SQL engine | | "Who is the founder of X?" | NL2SPARQL engine (KG) | | "Which quarter had highest sales?" | TableQA engine | | "Show average salary by department" | Auto-detect SQL | | Queries with aggregations, filters, joins | Route to SQL | | Entity relationship queries | Route to KG | | Questions about CSV/Excel files | Route to TableQA | | Multi-hop queries across sources | Decompose + fuse | ## Architecture ### 4-Layer Source Detection ```yaml Layer 1 (15%): Rule-Based - 20+ keywords per source type - 7 regex patterns (aggregation, comparison, relation) - Fast pre-filtering Layer 2 (35%): LLM Semantic - Intent classification - Entity/predicate detection - Multi-hop identification Layer 3a (25%): SQL Schema Match - Inverted index on tables/columns - Automatic JOIN inference - Confidence scoring Layer 3b (25%): KG Entity Link - Entity mention extraction - SPARQL endpoint lookup - Predicate pattern matching Layer 3c (25%+30%): Entity Verification - Cross-source entity existence check - 30% score boost for verified entities Layer 4: Multi-Source Fusion - Weighted aggregation - Execution plan generation ``` ### Query Generation Pipeline ``` 1. Schema/Entity Linking → Identify relevant tables/columns/entities 2. Parallel Generation → Generate 3 candidates concurrently 3. Multi-Round Revision → 2 rounds of self-review 4. Validation → Syntax and semantic checks 5. Voting → Select best candidate 6. Execution → Run query 7. Result Verification → Validate reasonableness ``` ## Engines ### NL2SQL Engine ```python from src.engines.nl2sql.multi_stage_engine import MultiStageNL2SQLEngine engine = MultiStageNL2SQLEngine({ "name": "sql_engine", "schema": schema, "llm_config": { "model": "deepseek-chat", "api_key": "sk-...", }, "generation_config": { "num_candidates": 3, "max_revisions": 2, "parallel_generation": True, }, }) result = await engine.execute("How many employees in Engineering?", {}) ``` **Features:** - Schema linking (rule-based + LLM) - Parallel SQL candidate generation - Multi-round self-revision - Voting mechanism - Result verification ### NL2SPARQL Engine ```python from src.engines.nl2sparql.multi_stage_engine import MultiStageNL2SPARQLEngine engine = MultiStageNL2SPARQLEngine({ "name": "sparql_engine", "endpoint_url": "https://dbpedia.org/sparql", "ontology": ontology, "llm_config": {"model": "gpt-4", "api_key": "sk-..."}, }) result = await engine.execute("Who founded Microsoft?", {}) ``` **Features:** - Entity linking to KG - Ontology retrieval - SPARQL generation with revision - Multi-endpoint support ### TableQA Engine ```python from src.engines.table_qa.multi_stage_engine import MultiStageTableQAEngine engine = MultiStageTableQAEngine({ "name": "table_engine", "table_path": "data/sales.csv", "llm_config": {"model": "deepseek-chat", "api_key": "sk-..."}, }) result = await engine.execute("Which quarter had highest sales?", {}) ``` **Features:** - Table schema analysis - Query intent interpretation - Pandas code generation - Safe execution sandbox ## Multi-LLM Support Override model and API key at runtime: ```python # Initialize with default engine = MultiStageNL2SQLEngine({ "llm_config": {"model": "deepseek-chat", "api_key": "sk-deepseek-key"}, }) # Override per-call result = await engine.execute( query="Complex query", context={}, model="gpt-4-turbo", # Override model [REDACTED_SECRET]", # Override API key ) ``` ### Supported Providers | Provider | Models | Configuration | |----------|--------|---------------| | DeepSeek | deepseek-chat | `base_url: https://api.deepseek.com/v1` | | OpenAI | gpt-4, gpt-3.5-turbo | Default endpoint | | Azure OpenAI | gpt-4 | `base_url: https://{resource}.openai.azure.com` | | Local (Ollama) | llama2, mistral | `base_url: http://localhost:11434/v1` | ## Configuration ### LLM Configuration ```yaml llm_config: model: deepseek-chat [REDACTED_SECRET] base_url: https://api.deepseek.com/v1 # Optional temperature: 0.1 max_tokens: 500 timeout: 30 ``` ### Generation Configuration ```yaml generation_config: num_candidates: 3 # SQL/SPARQL candidates to generate max_revisions: 2 # Self-revision rounds parallel_generation: true # Concurrent candidate generation voting_enabled: true # Multi-candidate voting ``` ### Source Detection Weights ```yaml weights: rule_based: 0.15 # Layer 1 llm_based: 0.35 # Layer 2 schema_based: 0.25 # Layer 3a/3b verification: 0.25 # Layer 3c verification_boost: 0.3 # 30% boost for verified entities ``` ## Workflows ### Complete Query Flow ```python from src.orchestrator import HeteroMindOrchestrator orchestrator = HeteroMindOrchestrator({ "source_detection": { "layer2": {"api_key": "sk-...", "model": "gpt-4"}, "layer3": {"schemas": [schema], "kg_endpoints": [...]}, }, "engines": { "sql": [{"name": "default", "enabled": True}], "sparql": [{"name": "default", "enabled": True}], "table_qa": [{"name": "default", "enabled": True}], }, }) response = await orchestrator.query("How many employees in Engineering?") print(f"Answer: {response.answer}") print(f"Source: {response.sources}") print(f"Confidence: {response.confidence:.2f}") ``` ### Source Detection Only ```python from src.classifier import SourceDetectorOrchestrator detector = SourceDetectorOrchestrator({ "layer2": {"api_key": "sk-...", "model": "gpt-4"}, "layer3": {"schemas": [schema]}, }) decision = await detector.detect("How many employees?") print(f"Primary Source: {decision.primary_source.value}") print(f"Confidence: {decision.confidence:.2f}") print(f"Execution Plan: {decision.execution_plan}") ``` ## Test Results | Engine | Tests | Passed | Accuracy | Avg Confidence | Avg Time | |--------|-------|--------|----------|----------------|----------| | SQL (NL2SQL) | 3 | 3 | 100.0% | 0.60 | 22.5s | | SPARQL (NL2SPARQL) | 2 | 2 | 100.0% | 0.20 | 36.3s | | TableQA | 3 | 3 | 100.0% | 0.62 | 24.2s | | **Overall** | 8 | 8 | 100.0% | 0.51 | 26.6s | ## Environment Variables ### Required (for LLM-based generation) | Variable | Description | Example | |----------|-------------|---------| | `DEEPSEEK_API_KEY` | DeepSeek API key | `sk-...` | | `OPENAI_API_KEY` | OpenAI API key | `sk-...` | ### Optional (for specific features) | Variable | Description | Example | |----------|-------------|---------| | `MYSQL_CONNECTION_STRING` | MySQL database connection | `mysql://user:pass@host/db` | | `CUSTOM_KG_ENDPOINT` | Custom KG SPARQL endpoint | `https://example.com/sparql` | | `WORKSPACE` | Base path for table file scanning | `/path/to/workspace` | ### Setup ```bash # Copy example env file cp .env.example .env # Edit with your credentials nano .env # Load environment export $(cat .env | xargs) ``` ## Installation ```bash cd HeteroMind pip install -r requirements.txt ``` ### Requirements - Python 3.10+ - `aiohttp`, `pandas`, `openpyxl` - OpenAI-compatible API key (optional) ## Project Structure ``` HeteroMind/ ├── src/ │ ├── classifier/ # 4-layer source detection │ │ ├── rule_detector.py # Layer 1 │ │ ├── llm_detector.py # Layer 2 │ │ ├── sql_schema_matcher.py # Layer 3a │ │ ├── kg_entity_linker.py # Layer 3b │ │ ├── entity_verifier.py # Layer 3c │ │ └── source_fusion.py # Layer 4 │ ├── engines/ # Query engines │ │ ├── nl2sql/ │ │ ├── nl2sparql/ │ │ └── table_qa/ │ ├── decomposer/ # Task decomposition │ ├── fusion/ # Result fusion │ ├── generator/ # Answer generation │ └── orchestrator.py # Main orchestrator ├── config/ │ └── source_detection.yaml ├── tests/ │ └── test_data/ ├── comprehensive_tests.py └── SKILL.md ``` ## Examples ### SQL: Aggregation with Filter **Query:** "How many employees are in the Engineering department?" **Generated SQL:** ```sql SELECT COUNT(*) FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'Engineering' ``` ### SPARQL: Entity Relationship **Query:** "Who is the founder of Microsoft?" **Generated SPARQL:** ```sparql SELECT ?founder WHERE { <http://dbpedia.org/resource/Microsoft> <http://dbpedia.org/ontology/founder> ?founder } ``` ### TableQA: Aggregation **Query:** "Which quarter had the highest sales in 2024?" **Generated Code:** ```python result = df.groupby('quarter')['sales'].sum().idxmax() ``` ## Skill Contract Skills that use HeteroMind should declare: ```yaml heteromind: reads: [Database Schema, KG Ontology, Table Files] writes: [Generated SQL, SPARQL, Pandas Code] requires: - LLM API key (for generation stages) - Schema metadata (for source detection) postconditions: - Generated query passes validation - Result verified for reasonableness ``` ## Integration Patterns ### With Agent Memory Log query execution for audit: ```python from src.orchestrator import HeteroMindOrchestrator orchestrator = HeteroMindOrchestrator(config) response = await orchestrator.query(query) # Log to agent memory memory.record({ "action": "knowledge_query", "query": query, "source": response.sources, "confidence": response.confidence, "answer": response.answer, }) ``` ### Multi-Source Fusion For queries requiring multiple sources: ```python # Query automatically detects hybrid need response = await orchestrator.query( "Show employees who published papers" ) # Routes to: SQL (employees) + KG (papers) + Fusion ``` ## References - `README.md` — Full documentation and API reference - `USAGE.md` — Detailed usage guide with multi-LLM examples - `config/source_detection.yaml` — Detection configuration - `tests/test_data/` — Example schemas and test data --- *Version: 0.1.0* *Last Updated: 2026-04-12* *Test Coverage: 100.0% accuracy on 8 test cases*

path: requirements.txt
content: # HeteroMind - Python Dependencies

# Core
pyyaml>=6.0
pydantic>=2.0

# Async support
aiohttp>=3.9.0

# LLM Integration (optional - choose your provider)
openai>=1.0.0
# anthropic>=0.18.0  # For Claude
# azure-ai-openai>=1.0.0  # For Azure OpenAI

# SQL Engines
sqlparse>=0.4.0
sqlalchemy>=2.0.0
# git+https://github.com/HKUSTDial/DeepEye-SQL.git  # DeepEye-SQL (when available)

# SPARQL Engines
rdflib>=6.0.0
sparqlwrapper>=2.0.0

# Table QA
pandas>=2.0.0
pandasai>=2.0.0
duckdb>=0.9.0
openpyxl>=3.1.0  # For Excel support

# Testing
pytest>=7.4.0
pytest-asyncio>=0.21.0

# Utilities
python-dotenv>=1.0.0  # For environment variables
logging>=0.5.1


path: config/source_detection.yaml
content: # HeteroMind - Source Detection Configuration
# Configures the 4-layer knowledge source detection system

# =============================================================================
# Layer Weights
# Controls contribution of each layer to the final decision
# All weights should sum to 1.0
# =============================================================================
weights:
  rule_based: 0.15      # Layer 1: Fast keyword/pattern matching
  llm_based: 0.35       # Layer 2: LLM semantic classification
  schema_based: 0.25    # Layer 3a/3b: Schema matching & entity linking
  verification: 0.25    # Layer 3c: Entity retrieval verification

# =============================================================================
# Confidence Thresholds
# Determines action based on final confidence score
# =============================================================================
thresholds:
  high_confidence: 0.95   # Auto-execute only with very high confidence
  medium_confidence: 0.7  # Execute with fallback strategies ready
  low_confidence: 0.5     # Ask user for clarification before proceeding

# Auto-execution control (SAFETY)
auto_execute: false         # Set to true only after testing
require_confirmation: true  # Always require user confirmation for execution

# =============================================================================
# Verification Boost
# Score boost applied when entities are verified to exist in source
# =============================================================================
verification_boost: 0.3   # 30% boost for verified entities

# =============================================================================
# Layer 2: LLM Configuration
# =============================================================================
llm:
  provider: "openai"      # openai, anthropic, azure, local
  model: "gpt-4"          # Model name
  temperature: 0.1        # Low temp for deterministic output
  max_tokens: 500         # Max response tokens
  timeout_seconds: 30     # Request timeout

# =============================================================================
# Layer 3: Schema & Knowledge Graph Configuration
# =============================================================================
schemas:
  # SQL Database schemas (auto-discovered or manually defined)
  sql_databases:
    - name: "company_db"
      type: "mysql"
      connection_env: "MYSQL_CONNECTION_STRING"
      auto_discover: true
      cache_ttl_seconds: 3600
      tables:
        - name: "employees"
          description: "Employee records with department, salary, hire date"
          columns:
            - name: "employee_id"
              type: "INT"
              alias: "emp_id"
            - name: "name"
              type: "VARCHAR"
            - name: "department_id"
              type: "INT"
            - name: "salary"
              type: "DECIMAL"
        - name: "projects"
          description: "Project records with budget, status, timeline"
          columns:
            - name: "project_id"
              type: "INT"
            - name: "name"
              type: "VARCHAR"
            - name: "budget"
              type: "DECIMAL"
            - name: "status"
              type: "VARCHAR"

knowledge_graph:
  # SPARQL endpoints
  endpoints:
    - name: "dbpedia"
      url: "https://dbpedia.org/sparql"
      enabled: true
      priority: 1
      timeout_seconds: 10
    
    - name: "wikidata"
      url: "https://query.wikidata.org/sparql"
      enabled: true
      priority: 2
      timeout_seconds: 10
    
    - name: "custom_kg"
      url_env: "CUSTOM_KG_ENDPOINT"
      enabled: false
      priority: 3
  
  # Entity linking settings
  entity_linking:
    max_candidates: 5     # Max candidates per entity mention
    min_similarity: 0.7   # Min similarity score for matches
    use_fuzzy: true       # Enable fuzzy matching
    fuzzy_threshold: 0.8  # Fuzzy match threshold

# =============================================================================
# Layer 3c: Entity Verification Configuration
# =============================================================================
entity_verification:
  # SQL verification
  sql:
    enabled: true
    search_columns: ["name", "title", "label", "description"]
    case_sensitive: false
    min_match_length: 3
  
  # KG verification
  kg:
    enabled: true
    use_ask_query: true   # Use SPARQL ASK for existence check
    cache_results: true
    cache_ttl_seconds: 300
  
  # Table file verification (DISABLED BY DEFAULT FOR SAFETY)
  tables:
    enabled: false          # Set to true only after explicit user consent
    paths: []               # Manually specify paths, no wildcards
    search_all_columns: false
    max_file_size_mb: 10
    explicit_paths_only: true  # Only scan explicitly specified files

# =============================================================================
# Execution Plan Settings
# =============================================================================
execution:
  # Parallel execution settings
  parallel:
    enabled: true
    max_concurrent: 3     # Max sources to query in parallel
  
  # Fallback strategy
  fallback:
    enabled: true
    max_retries: 2
    retry_delay_seconds: 1
  
  # Timeout settings
  timeout:
    total_seconds: 60     # Total timeout for all sources
    per_source_seconds: 20

# =============================================================================
# Logging & Debugging
# =============================================================================
logging:
  level: "INFO"           # DEBUG, INFO, WARNING, ERROR
  log_layer_outputs: true
  log_verification_details: true
  
  # Debug mode: include all intermediate results in final decision
  debug_mode: false

# =============================================================================
# Performance Tuning
# =============================================================================
performance:
  # Cache settings
  cache:
    enabled: true
    ttl_seconds: 300      # Cache TTL for detection results
    max_size: 1000        # Max cached queries
  
  # Batch settings
  batch:
    enabled: true
    max_batch_size: 10    # Max entities to verify in one batch
  
  # Async settings
  async:
    enabled: true
    max_concurrent_tasks: 5


path: USAGE.md
content: # HeteroMind Usage Guide

## Quick Start

### Basic Usage with Default LLM

```python
import asyncio
from src.engines.nl2sql.multi_stage_engine import MultiStageNL2SQLEngine

# Initialize with default LLM configuration
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "deepseek-chat",
        "api_key": "sk-your-api-key",
        "base_url": "https://api.deepseek.com/v1",
    },
})

# Execute query
result = await engine.execute("How many employees in Engineering?", {})
print(f"Generated SQL: {result.generated_sql}")
```

### Using Different LLM Providers

#### DeepSeek

```python
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "deepseek-chat",
        "api_key": "sk-...",
        "base_url": "https://api.deepseek.com/v1",
    },
})
```

#### OpenAI GPT-4

```python
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "gpt-4",
        "api_key": "sk-...",
    },
})
```

#### Azure OpenAI

```python
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "gpt-4",
        "api_key": "your-azure-key",
        "base_url": "https://your-resource.openai.azure.com/v1",
    },
})
```

#### Local LLM (Ollama)

```python
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "llama2",
        "base_url": "http://localhost:11434/v1",
        "api_key": "not-needed",  # Ollama doesn't require API key
    },
})
```

### Per-Call LLM Override

Override model and API key at runtime without reinitializing the engine:

```python
# Initialize with default LLM
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "deepseek-chat",
        "api_key": "sk-deepseek-key",
    },
})

# Use default LLM
result1 = await engine.execute("Query 1", {})

# Override with GPT-4 for specific query
result2 = await engine.execute(
    query="Query 2",
    context={},
    model="gpt-4-turbo",
    [REDACTED_SECRET]",
)

# Override with different model for another query
result3 = await engine.execute(
    query="Query 3",
    context={},
    model="claude-3-sonnet",
    [REDACTED_SECRET]",
)
```

### Multi-Engine Setup with Different LLMs

```python
from src.engines.nl2sql.multi_stage_engine import MultiStageNL2SQLEngine
from src.engines.nl2sparql.multi_stage_engine import MultiStageNL2SPARQLEngine
from src.engines.table_qa.multi_stage_engine import MultiStageTableQAEngine

# SQL Engine with DeepSeek
sql_engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": sql_schema,
    "llm_config": {
        "model": "deepseek-chat",
        "api_key": "sk-deepseek-key",
    },
})

# SPARQL Engine with GPT-4
sparql_engine = MultiStageNL2SPARQLEngine({
    "name": "sparql_engine",
    "ontology": kg_ontology,
    "llm_config": {
        "model": "gpt-4",
        "api_key": "sk-openai-key",
    },
})

# TableQA Engine with local LLM
table_engine = MultiStageTableQAEngine({
    "name": "table_engine",
    "table_path": "data.csv",
    "llm_config": {
        "model": "llama2",
        "base_url": "http://localhost:11434/v1",
    },
})
```

## Configuration Options

### LLM Configuration Parameters

| Parameter | Required | Description | Example |
|-----------|----------|-------------|---------|
| `model` | Yes | Model name | `"gpt-4"`, `"deepseek-chat"`, `"claude-3"` |
| `api_key` | Yes* | API key for authentication | `"sk-..."` |
| `base_url` | No | Custom API endpoint | `"https://api.deepseek.com/v1"` |
| `temperature` | No | Sampling temperature (0.0-1.0) | `0.1` for deterministic |
| `max_tokens` | No | Max response tokens | `500` |
| `timeout` | No | Request timeout in seconds | `30` |

*Required for cloud providers, optional for local LLMs

### Generation Configuration Parameters

| Parameter | Default | Description |
|-----------|---------|-------------|
| `num_candidates` | 3 | Number of SQL/SPARQL candidates to generate |
| `max_revisions` | 2 | Maximum self-revision rounds |
| `parallel_generation` | True | Enable parallel candidate generation |
| `voting_enabled` | True | Enable multi-candidate voting |

## Best Practices

### 1. Choose the Right Model

- **Simple queries**: Use cost-effective models like `deepseek-chat` or `gpt-3.5-turbo`
- **Complex queries**: Use powerful models like `gpt-4` or `claude-3-opus`
- **Schema linking**: Can use smaller models (high pattern matching)
- **Query generation**: Use larger models for better accuracy

### 2. API Key Management

```python
import os

# Load from environment variable
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "gpt-4",
        "api_key": os.getenv("OPENAI_API_KEY"),
    },
})
```

### 3. Cost Optimization

```python
# Use cheaper model for most queries
default_engine = MultiStageNL2SQLEngine({
    "llm_config": {"model": "deepseek-chat", "api_key": "..."},
})

# Override with expensive model only for complex queries
result = await default_engine.execute(
    complex_query,
    {},
    model="gpt-4",
    api_key="...",
)
```

### 4. Error Handling

```python
try:
    result = await engine.execute(query, {}, model="gpt-4", api_key="...")
    if result.success:
        print(f"Success: {result.generated_sql}")
    else:
        print(f"Error: {result.error}")
except Exception as e:
    print(f"Exception: {e}")
```

## Troubleshooting

### Invalid API Key

```
Error: 401 Unauthorized
```

**Solution**: Check that `api_key` is correct and has sufficient credits.

### Model Not Found

```
Error: 404 Model not found
```

**Solution**: Verify `model` name is correct for the provider.

### Rate Limiting

```
Error: 429 Too Many Requests
```

**Solution**: 
- Reduce `num_candidates`
- Add retry logic with backoff
- Use `timeout` parameter

### Local LLM Connection Failed

```
Error: Connection refused
```

**Solution**: 
- Ensure Ollama/local LLM server is running
- Check `base_url` is correct
- Verify port is not blocked by firewall

## Examples

### Example 1: Multi-Provider Fallback

```python
async def execute_with_fallback(query, schema):
    providers = [
        {"model": "gpt-4", "api_key": os.getenv("OPENAI_API_KEY")},
        {"model": "deepseek-chat", "api_key": os.getenv("DEEPSEEK_API_KEY")},
        {"model": "llama2", "base_url": "http://localhost:11434/v1"},
    ]
    
    engine = MultiStageNL2SQLEngine({
        "name": "sql_engine",
        "schema": schema,
    })
    
    for provider in providers:
        try:
            result = await engine.execute(query, {}, **provider)
            if result.success:
                return result
        except Exception as e:
            print(f"Provider {provider['model']} failed: {e}")
    
    raise Exception("All providers failed")
```

### Example 2: Batch Processing with Different Models

```python
async def process_queries(queries, schema):
    engine = MultiStageNL2SQLEngine({
        "name": "sql_engine",
        "schema": schema,
        "llm_config": {"model": "deepseek-chat", "api_key": "..."},
    })
    
    results = []
    for query in queries:
        # Use GPT-4 for complex queries
        if len(query.split()) > 15:
            result = await engine.execute(
                query, {},
                model="gpt-4",
                api_key="...",
            )
        else:
            result = await engine.execute(query, {})
        
        results.append(result)
    
    return results
```

---

For more information, see [README.md](README.md).


path: tests/test_data/company_db_schema.json
content: {
  "database": "company_db",
  "description": "Company employee and project management database",
  "tables": [
    {
      "name": "employees",
      "description": "Employee information",
      "columns": [
        {"name": "id", "type": "INT", "primary_key": true, "description": "Employee ID"},
        {"name": "name", "type": "VARCHAR(100)", "description": "Employee name"},
        {"name": "email", "type": "VARCHAR(100)", "description": "Email address"},
        {"name": "department_id", "type": "INT", "foreign_key": {"table": "departments", "column": "id"}, "description": "Department ID"},
        {"name": "salary", "type": "DECIMAL(10,2)", "description": "Annual salary"},
        {"name": "hire_date", "type": "DATE", "description": "Date of hire"},
        {"name": "title", "type": "VARCHAR(50)", "description": "Job title"}
      ]
    },
    {
      "name": "departments",
      "description": "Department information",
      "columns": [
        {"name": "id", "type": "INT", "primary_key": true, "description": "Department ID"},
        {"name": "name", "type": "VARCHAR(100)", "description": "Department name"},
        {"name": "location", "type": "VARCHAR(100)", "description": "Office location"},
        {"name": "budget", "type": "DECIMAL(12,2)", "description": "Annual department budget"}
      ]
    },
    {
      "name": "projects",
      "description": "Project information",
      "columns": [
        {"name": "id", "type": "INT", "primary_key": true, "description": "Project ID"},
        {"name": "name", "type": "VARCHAR(200)", "description": "Project name"},
        {"name": "department_id", "type": "INT", "foreign_key": {"table": "departments", "column": "id"}, "description": "Owning department"},
        {"name": "budget", "type": "DECIMAL(12,2)", "description": "Project budget"},
        {"name": "start_date", "type": "DATE", "description": "Project start date"},
        {"name": "status", "type": "VARCHAR(20)", "description": "Project status: active, completed, on_hold"}
      ]
    },
    {
      "name": "employee_projects",
      "description": "Employee project assignments",
      "columns": [
        {"name": "employee_id", "type": "INT", "foreign_key": {"table": "employees", "column": "id"}, "description": "Employee ID"},
        {"name": "project_id", "type": "INT", "foreign_key": {"table": "projects", "column": "id"}, "description": "Project ID"},
        {"name": "role", "type": "VARCHAR(50)", "description": "Role in project"},
        {"name": "allocation", "type": "DECIMAL(3,2)", "description": "Time allocation (0.0-1.0)"}
      ]
    }
  ],
  "relationships": [
    {
      "from_table": "employees",
      "from_column": "department_id",
      "to_table": "departments",
      "to_column": "id",
      "type": "many_to_one"
    },
    {
      "from_table": "projects",
      "from_column": "department_id",
      "to_table": "departments",
      "to_column": "id",
      "type": "many_to_one"
    },
    {
      "from_table": "employee_projects",
      "from_column": "employee_id",
      "to_table": "employees",
      "to_column": "id",
      "type": "many_to_one"
    },
    {
      "from_table": "employee_projects",
      "from_column": "project_id",
      "to_table": "projects",
      "to_column": "id",
      "type": "many_to_one"
    }
  ]
}


path: tests/test_data/sales_quarterly.csv
content: quarter,year,region,product_category,sales,orders,customers,returns
Q1,2024,North,Electronics,1250.5,450,320,15
Q1,2024,North,Clothing,890.3,520,410,28
Q1,2024,South,Electronics,980.2,380,290,12
Q1,2024,South,Clothing,720.8,430,350,22
Q1,2024,East,Electronics,1100.0,420,310,18
Q1,2024,East,Clothing,650.5,380,290,19
Q1,2024,West,Electronics,1350.8,480,350,20
Q1,2024,West,Clothing,780.2,450,360,25
Q2,2024,North,Electronics,1380.2,490,350,18
Q2,2024,North,Clothing,920.5,540,430,30
Q2,2024,South,Electronics,1050.8,400,310,14
Q2,2024,South,Clothing,780.3,460,380,24
Q2,2024,East,Electronics,1200.5,450,340,20
Q2,2024,East,Clothing,700.2,400,310,21
Q2,2024,West,Electronics,1450.0,510,380,22
Q2,2024,West,Clothing,820.8,470,380,27
Q3,2024,North,Electronics,1520.8,530,390,20
Q3,2024,North,Clothing,1050.2,590,470,32
Q3,2024,South,Electronics,1180.5,440,340,16
Q3,2024,South,Clothing,850.0,490,400,26
Q3,2024,East,Electronics,1350.2,490,370,22
Q3,2024,East,Clothing,780.5,440,340,23
Q3,2024,West,Electronics,1580.0,550,410,24
Q3,2024,West,Clothing,900.5,510,410,29
Q4,2024,North,Electronics,2100.5,720,520,28
Q4,2024,North,Clothing,1450.8,820,650,42
Q4,2024,South,Electronics,1680.2,610,470,22
Q4,2024,South,Clothing,1180.5,670,540,35
Q4,2024,East,Electronics,1850.0,670,510,30
Q4,2024,East,Clothing,1100.2,620,490,33
Q4,2024,West,Electronics,2250.8,780,580,32
Q4,2024,West,Clothing,1520.0,860,680,45


path: tests/test_data/kg_ontology.json
content: {
  "name": "Research Knowledge Graph",
  "description": "Ontology for academic publications and researcher information",
  "prefixes": {
    "foaf": "http://xmlns.com/foaf/0.1/",
    "dc": "http://purl.org/dc/terms/",
    "dbo": "http://dbpedia.org/ontology/",
    "dbp": "http://dbpedia.org/property/",
    "schema": "http://schema.org/"
  },
  "classes": [
    {
      "name": "Person",
      "uri": "http://xmlns.com/foaf/0.1/Person",
      "labels": ["Person", "人", "Researcher", "研究者"],
      "synonyms": ["Individual", "Author", "科学家", "学者"]
    },
    {
      "name": "Organization",
      "uri": "http://schema.org/Organization",
      "labels": ["Organization", "机构", "Company", "公司", "University", "大学"]
    },
    {
      "name": "Paper",
      "uri": "http://schema.org/ScholarlyArticle",
      "labels": ["Paper", "论文", "Publication", "出版物", "Article", "文章"]
    },
    {
      "name": "Conference",
      "uri": "http://schema.org/Event",
      "labels": ["Conference", "会议", "Workshop", "研讨会"]
    },
    {
      "name": "ResearchField",
      "uri": "http://dbpedia.org/ontology/AcademicSubject",
      "labels": ["Research Field", "研究领域", "Topic", "主题"]
    }
  ],
  "properties": [
    {
      "name": "name",
      "uri": "http://xmlns.com/foaf/0.1/name",
      "labels": ["name", "姓名", "title"],
      "domain": "Person",
      "range": "string"
    },
    {
      "name": "founder",
      "uri": "http://dbpedia.org/ontology/founder",
      "labels": ["founder", "创始人", "founded by", "创立于"],
      "domain": "Organization",
      "range": "Person"
    },
    {
      "name": "ceo",
      "uri": "http://dbpedia.org/ontology/keyPerson",
      "labels": ["CEO", "首席执行官", "leader", "领导者"],
      "domain": "Organization",
      "range": "Person"
    },
    {
      "name": "author",
      "uri": "http://purl.org/dc/terms/creator",
      "labels": ["author", "作者", "written by", "撰写者"],
      "domain": "Paper",
      "range": "Person"
    },
    {
      "name": "publishedIn",
      "uri": "http://purl.org/dc/terms/isPartOf",
      "labels": ["published in", "发表于", "venue"],
      "domain": "Paper",
      "range": "Conference"
    },
    {
      "name": "researchField",
      "uri": "http://purl.org/dc/terms/subject",
      "labels": ["research field", "研究领域", "topic", "主题", "about"],
      "domain": "Paper",
      "range": "ResearchField"
    },
    {
      "name": "affiliation",
      "uri": "http://schema.org/affiliation",
      "labels": ["affiliation", "所属机构", "works at", "工作于"],
      "domain": "Person",
      "range": "Organization"
    },
    {
      "name": "citationCount",
      "uri": "http://schema.org/citationCount",
      "labels": ["citation count", "引用次数", "citations"],
      "domain": "Paper",
      "range": "integer"
    }
  ],
  "relationships": [
    {
      "predicate": "http://dbpedia.org/ontology/founder",
      "domain": "http://schema.org/Organization",
      "range": "http://xmlns.com/foaf/0.1/Person",
      "description": "Links an organization to its founder(s)"
    },
    {
      "predicate": "http://purl.org/dc/terms/creator",
      "domain": "http://schema.org/ScholarlyArticle",
      "range": "http://xmlns.com/foaf/0.1/Person",
      "description": "Links a paper to its author(s)"
    },
    {
      "predicate": "http://purl.org/dc/terms/subject",
      "domain": "http://schema.org/ScholarlyArticle",
      "range": "http://dbpedia.org/ontology/AcademicSubject",
      "description": "Links a paper to its research field/topic"
    }
  ],
  "sample_entities": {
    "persons": [
      {"uri": "http://dbpedia.org/resource/Bill_Gates", "label": "Bill Gates", "labels_zh": "比尔·盖茨"},
      {"uri": "http://dbpedia.org/resource/Paul_Allen", "label": "Paul Allen", "labels_zh": "保罗·艾伦"},
      {"uri": "http://dbpedia.org/resource/Steve_Jobs", "label": "Steve Jobs", "labels_zh": "史蒂夫·乔布斯"},
      {"uri": "http://dbpedia.org/resource/Elon_Musk", "label": "Elon Musk", "labels_zh": "埃隆·马斯克"}
    ],
    "organizations": [
      {"uri": "http://dbpedia.org/resource/Microsoft", "label": "Microsoft", "labels_zh": "微软"},
      {"uri": "http://dbpedia.org/resource/Apple_Inc.", "label": "Apple Inc.", "labels_zh": "苹果公司"},
      {"uri": "http://dbpedia.org/resource/Tesla,_Inc.", "label": "Tesla, Inc.", "labels_zh": "特斯拉"}
    ],
    "research_fields": [
      {"uri": "http://dbpedia.org/resource/Artificial_intelligence", "label": "Artificial Intelligence", "labels_zh": "人工智能"},
      {"uri": "http://dbpedia.org/resource/Machine_learning", "label": "Machine Learning", "labels_zh": "机器学习"},
      {"uri": "http://dbpedia.org/resource/Natural_language_processing", "label": "Natural Language Processing", "labels_zh": "自然语言处理"}
    ]
  }
}


path: tests/comprehensive_test_report.md
content: # HeteroMind Comprehensive Test Report

**Generated:** 2026-04-12 17:42:50
**Model:** deepseek-chat

## Summary

| Metric | Value |
|--------|-------|
| Total Tests | 8 |
| Passed | 8 ✅ |
| Failed | 0 ❌ |
| **Overall Accuracy** | **100.0%** |
| Avg Confidence | 0.64 |
| Avg Response Time | 25246ms |

## Results by Engine

| Engine | Tests | Passed | Accuracy | Avg Confidence | Avg Time |
|--------|-------|--------|----------|----------------|----------|
| SQL | 3 | 3 | 100.0% | 0.60 | 21421ms |
| SPARQL | 2 | 2 | 100.0% | 0.72 | 37752ms |
| TableQA | 3 | 3 | 100.0% | 0.62 | 20735ms |

## Detailed Results

### SQL-001 - ✅ PASS

**Query:** How many employees are in the Engineering department?

**Engine:** SQL

**Confidence:** 0.60

**Execution Time:** 25618ms

**Generated Query:**
```
SELECT COUNT(*) 
FROM employees e
JOIN departments d ON e.department_id = d.id
WHERE d.name = 'Engineering';
```

### SQL-002 - ✅ PASS

**Query:** 按部门统计每个部门的平均薪资，并显示薪资最高的前 3 个部门

**Engine:** SQL

**Confidence:** 0.60

**Execution Time:** 22919ms

**Generated Query:**
```
SELECT 
    d.name AS department_name,
    AVG(e.salary) AS average_salary
FROM 
    departments d
JOIN 
    employees e ON d.id = e.department_id
GROUP BY 
    d.id, d.name
ORDER BY 
    average_salary DESC
LIMIT 3;
```

### SQL-003 - ✅ PASS

**Query:** List all employees hired after 2020

**Engine:** SQL

**Confidence:** 0.60

**Execution Time:** 15727ms

**Generated Query:**
```
SELECT * FROM employees WHERE hire_date > '2020-12-31';
```

### KG-001 - ✅ PASS

**Query:** Who is the founder of Microsoft?

**Engine:** SPARQL

**Confidence:** 0.60

**Execution Time:** 27886ms

**Generated Query:**
```
PREFIX dbo: <http://dbpedia.org/ontology/>
PREFIX dbr: <http://dbpedia.org/resource/>

SELECT ?founder WHERE {
    dbr:Microsoft dbo:founder ?founder .
}
```

### KG-002 - ✅ PASS

**Query:** What companies are headquartered in Seattle?

**Engine:** SPARQL

**Confidence:** 0.85

**Execution Time:** 47618ms

**Generated Query:**
```
PREFIX dbo: <http://dbpedia.org/ontology/>
PREFIX dbr: <http://dbpedia.org/resource/>
PREFIX rdf: <http://www.w3.org/1999/02/22-rdf-syntax-ns#>
PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>

SELECT DISTINCT ?company ?companyLabel WHERE {
  ?company rdf:type dbo:Company ;
           dbo:headquarter dbr:Seattle ;
           rdfs:label ?companyLabel .
  FILTER(LANG(?companyLabel) = "en")
}
```

### TQ-001 - ✅ PASS

**Query:** 2024 年哪个季度的销售额最高？

**Engine:** TableQA

**Confidence:** 0.50

**Execution Time:** 25327ms

**Generated Query:**
```
import pandas as pd

# Filter data for year 2024
df_2024 = df[df['year'] == 2024]

# Group by quarter and calculate total sales per quarter
quarterly_sales = df_2024.groupby('quarter')['sales'].sum().reset_index()

# Find the quarter with the maximum sales
max_sales_quarter = quarterly_sales.loc[quarterly_sales['sales'].idxmax()]

print(max_sales_quarter['quarter'])
```

### TQ-002 - ✅ PASS

**Query:** What is the average sales per quarter?

**Engine:** TableQA

**Confidence:** 0.85

**Execution Time:** 23538ms

**Generated Query:**
```
import pandas as pd

# Calculate average sales per quarter
result = df.groupby('quarter')['sales'].mean().reset_index()
result.columns = ['quarter', 'average_sales']
print(result)
```

### TQ-003-UNKNOWN [UNKNOWN SOURCE] - ✅ PASS

**Query:** 根据提供的表格，计算总销售额

**Engine:** TableQA

**Confidence:** 0.50

**Execution Time:** 13339ms

**Generated Query:**
```
total_sales = df['sales'].sum()
print(total_sales)
```



path: tests/README.md
content: # HeteroMind Test Suite

## Overview

This test suite validates the HeteroMind heterogeneous knowledge QA system across:
- **Multiple Languages**: Chinese (中文), English
- **Multiple Knowledge Sources**: SQL, SPARQL/KG, Table/CSV
- **Unknown Source Handling**: Tests panics with unspecified knowledge sources

## Test Cases Summary

| ID | Language | Source | Description |
|----|----------|--------|-------------|
| SQL-001 | English | SQL | Count employees in department |
| SQL-002 | Chinese | SQL | GROUP BY + AVG + ORDER BY + LIMIT |
| SQL-003 | English | SQL | Filter by date |
| KG-001 | English | SPARQL | Entity lookup (founder) |
| KG-002 | English | SPARQL | Entity lookup (headquarters) |
| TQ-001 | Chinese | Table | Find max sales by quarter |
| TQ-002 | English | Table | Calculate average |
| TQ-003-UNKNOWN | Chinese | Table | Unknown source test |

## Running Tests

### Prerequisites

```bash
pip install pandas openpyxl aiohttp
```

### Run All Tests (Recommended)

```bash
cd ~/Documents/openclaw/workspace/HeteroMind
python3 comprehensive_tests.py
```

This will:
- Test all 3 engines (SQL, SPARQL, TableQA)
- Run 8 test cases total
- Generate detailed reports:
  - `tests/comprehensive_test_results.json` - JSON format results
  - `tests/comprehensive_test_report.md` - Markdown format report

## Test Data

Located in `tests/test_data/`:

| File | Description |
|------|-------------|
| `company_db_schema.json` | SQL database schema |
| `sales_quarterly.csv` | Sample sales data (32 rows) |
| `kg_ontology.json` | Knowledge graph ontology |

## Evaluation Metrics

| Metric | Target | Description |
|--------|--------|-------------|
| Source Detection Accuracy | >90% | Correctly identifies knowledge source |
| Query Generation | >95% | Generates syntactically correct query |
| Execution Success | >90% | Query executes without error |
| Answer Accuracy | >80% | Answer matches expected result |

## Adding New Test Cases

Edit `comprehensive_tests.py` and add to the appropriate test case list:

```python
SQL_TEST_CASES = [
    {
        "id": "SQL-004",
        "query": "Your new query",
        "language": "en",
        "expected_sql": "Expected SQL pattern",
        "expected_tables": ["table1", "table2"],
    },
]
```

## Troubleshooting

### Test Fails with "No module named 'engines'"
```bash
# Ensure you're running from the HeteroMind directory
cd ~/Documents/openclaw/workspace/HeteroMind
```

### LLM Tests Fail Without API Key
```bash
# The test uses DeepSeek API by default
# Check DEEPSEEK_API_KEY in comprehensive_tests.py
```

## Contact

For questions or issues, please open a GitHub issue.


path: tests/comprehensive_test_results.json
content: {
  "timestamp": "2026-04-12T17:42:50.098321",
  "summary": {
    "total_tests": 8,
    "passed_tests": 8,
    "failed_tests": 0,
    "accuracy": 1.0,
    "avg_confidence": 0.6375,
    "avg_time_ms": 25246.495634317398
  },
  "by_engine": {
    "SQL": {
      "total": 3,
      "passed": 3,
      "accuracy": 1.0,
      "avg_confidence": 0.6,
      "avg_time_ms": 21421.34364446004
    },
    "SPARQL": {
      "total": 2,
      "passed": 2,
      "accuracy": 1.0,
      "avg_confidence": 0.725,
      "avg_time_ms": 37752.023577690125
    },
    "TableQA": {
      "total": 3,
      "passed": 3,
      "accuracy": 1.0,
      "avg_confidence": 0.6166666666666667,
      "avg_time_ms": 20734.6289952596
    }
  },
  "detailed_results": [
    {
      "test_id": "SQL-001",
      "query": "How many employees are in the Engineering department?",
      "engine_type": "SQL",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.6,
      "execution_time_ms": 25617.799043655396,
      "error": null,
      "generated_query": "SELECT COUNT(*) \nFROM employees e\nJOIN departments d ON e.department_id = d.id\nWHERE d.name = 'Engineering';",
      "details": {
        "schema_links": null,
        "revision_count": 3
      }
    },
    {
      "test_id": "SQL-002",
      "query": "按部门统计每个部门的平均薪资，并显示薪资最高的前 3 个部门",
      "engine_type": "SQL",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.6,
      "execution_time_ms": 22919.022798538208,
      "error": null,
      "generated_query": "SELECT \n    d.name AS department_name,\n    AVG(e.salary) AS average_salary\nFROM \n    departments d\nJOIN \n    employees e ON d.id = e.department_id\nGROUP BY \n    d.id, d.name\nORDER BY \n    average_salary DESC\nLIMIT 3;",
      "details": {
        "schema_links": null,
        "revision_count": 3
      }
    },
    {
      "test_id": "SQL-003",
      "query": "List all employees hired after 2020",
      "engine_type": "SQL",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.6,
      "execution_time_ms": 15727.209091186523,
      "error": null,
      "generated_query": "SELECT * FROM employees WHERE hire_date > '2020-12-31';",
      "details": {
        "schema_links": null,
        "revision_count": 3
      }
    },
    {
      "test_id": "KG-001",
      "query": "Who is the founder of Microsoft?",
      "engine_type": "SPARQL",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.6,
      "execution_time_ms": 27886.121034622192,
      "error": null,
      "generated_query": "PREFIX dbo: <http://dbpedia.org/ontology/>\nPREFIX dbr: <http://dbpedia.org/resource/>\n\nSELECT ?founder WHERE {\n    dbr:Microsoft dbo:founder ?founder .\n}",
      "details": {
        "linked_entities": "[LinkedEntity(mention='Microsoft', uri='http://dbpedia.org/resource/Microsoft', label='Microsoft', score=0.95, types=['Organization']), LinkedEntity(mention='founder', uri='http://dbpedia.org/ontology/founder', label='founder', score=0.9, types=['Property'])]",
        "revision_count": 3
      }
    },
    {
      "test_id": "KG-002",
      "query": "What companies are headquartered in Seattle?",
      "engine_type": "SPARQL",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.85,
      "execution_time_ms": 47617.92612075806,
      "error": null,
      "generated_query": "PREFIX dbo: <http://dbpedia.org/ontology/>\nPREFIX dbr: <http://dbpedia.org/resource/>\nPREFIX rdf: <http://www.w3.org/1999/02/22-rdf-syntax-ns#>\nPREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>\n\nSELECT DISTINCT ?company ?companyLabel WHERE {\n  ?company rdf:type dbo:Company ;\n           dbo:headquarter dbr:Seattle ;\n           rdfs:label ?companyLabel .\n  FILTER(LANG(?companyLabel) = \"en\")\n}",
      "details": {
        "linked_entities": "[LinkedEntity(mention='Seattle', uri='http://dbpedia.org/resource/Seattle', label='Seattle', score=0.9, types=['Organization'])]",
        "revision_count": 3
      }
    },
    {
      "test_id": "TQ-001",
      "query": "2024 年哪个季度的销售额最高？",
      "engine_type": "TableQA",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.5,
      "execution_time_ms": 25327.16417312622,
      "error": null,
      "generated_query": "import pandas as pd\n\n# Filter data for year 2024\ndf_2024 = df[df['year'] == 2024]\n\n# Group by quarter and calculate total sales per quarter\nquarterly_sales = df_2024.groupby('quarter')['sales'].sum().reset_index()\n\n# Find the quarter with the maximum sales\nmax_sales_quarter = quarterly_sales.loc[quarterly_sales['sales'].idxmax()]\n\nprint(max_sales_quarter['quarter'])",
      "details": {
        "table_schema": {
          "columns": [
            {
              "name": "quarter",
              "dtype": "object",
              "non_null": 32,
              "unique": 4,
              "sample_values": [
                "Q1",
                "Q1",
                "Q1"
              ]
            },
            {
              "name": "year",
              "dtype": "int64",
              "non_null": 32,
              "unique": 1,
              "sample_values": [
                2024,
                2024,
                2024
              ]
            },
            {
              "name": "region",
              "dtype": "object",
              "non_null": 32,
              "unique": 4,
              "sample_values": [
                "North",
                "North",
                "South"
              ]
            },
            {
              "name": "product_category",
              "dtype": "object",
              "non_null": 32,
              "unique": 2,
              "sample_values": [
                "Electronics",
                "Clothing",
                "Electronics"
              ]
            },
            {
              "name": "sales",
              "dtype": "float64",
              "non_null": 32,
              "unique": 31,
              "sample_values": [
                1250.5,
                890.3,
                980.2
              ]
            },
            {
              "name": "orders",
              "dtype": "int64",
              "non_null": 32,
              "unique": 23,
              "sample_values": [
                450,
                520,
                380
              ]
            },
            {
              "name": "customers",
              "dtype": "int64",
              "non_null": 32,
              "unique": 20,
              "sample_values": [
                320,
                410,
                290
              ]
            },
            {
              "name": "returns",
              "dtype": "int64",
              "non_null": 32,
              "unique": 22,
              "sample_values": [
                15,
                28,
                12
              ]
            }
          ],
          "row_count": 32,
          "dtypes": {
            "quarter": "object",
            "year": "int64",
            "region": "object",
            "product_category": "object",
            "sales": "float64",
            "orders": "int64",
            "customers": "int64",
            "returns": "int64"
          },
          "statistics": {
            "year": {
              "mean": 2024.0,
              "min": 2024.0,
              "max": 2024.0,
              "std": 0.0
            },
            "sales": {
              "mean": 1199.1343749999999,
              "min": 650.5,
              "max": 2250.8,
              "std": 404.6938011030187
            },
            "orders": {
              "mean": 531.875,
              "min": 380.0,
              "max": 860.0,
              "std": 127.29030827460966
            },
            "customers": {
              "mean": 410.3125,
              "min": 290.0,
              "max": 680.0,
              "std": 100.33763565598034
            },
            "returns": {
              "mean": 24.8125,
              "min": 12.0,
              "max": 45.0,
              "std": 7.570901425001964
            }
          }
        },
        "revision_count": 3,
        "is_unknown_source": false
      }
    },
    {
      "test_id": "TQ-002",
      "query": "What is the average sales per quarter?",
      "engine_type": "TableQA",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.85,
      "execution_time_ms": 23537.906885147095,
      "error": null,
      "generated_query": "import pandas as pd\n\n# Calculate average sales per quarter\nresult = df.groupby('quarter')['sales'].mean().reset_index()\nresult.columns = ['quarter', 'average_sales']\nprint(result)",
      "details": {
        "table_schema": {
          "columns": [
            {
              "name": "quarter",
              "dtype": "object",
              "non_null": 32,
              "unique": 4,
              "sample_values": [
                "Q1",
                "Q1",
                "Q1"
              ]
            },
            {
              "name": "year",
              "dtype": "int64",
              "non_null": 32,
              "unique": 1,
              "sample_values": [
                2024,
                2024,
                2024
              ]
            },
            {
              "name": "region",
              "dtype": "object",
              "non_null": 32,
              "unique": 4,
              "sample_values": [
                "North",
                "North",
                "South"
              ]
            },
            {
              "name": "product_category",
              "dtype": "object",
              "non_null": 32,
              "unique": 2,
              "sample_values": [
                "Electronics",
                "Clothing",
                "Electronics"
              ]
            },
            {
              "name": "sales",
              "dtype": "float64",
              "non_null": 32,
              "unique": 31,
              "sample_values": [
                1250.5,
                890.3,
                980.2
              ]
            },
            {
              "name": "orders",
              "dtype": "int64",
              "non_null": 32,
              "unique": 23,
              "sample_values": [
                450,
                520,
                380
              ]
            },
            {
              "name": "customers",
              "dtype": "int64",
              "non_null": 32,
              "unique": 20,
              "sample_values": [
                320,
                410,
                290
              ]
            },
            {
              "name": "returns",
              "dtype": "int64",
              "non_null": 32,
              "unique": 22,
              "sample_values": [
                15,
                28,
                12
              ]
            }
          ],
          "row_count": 32,
          "dtypes": {
            "quarter": "object",
            "year": "int64",
            "region": "object",
            "product_category": "object",
            "sales": "float64",
            "orders": "int64",
            "customers": "int64",
            "returns": "int64"
          },
          "statistics": {
            "year": {
              "mean": 2024.0,
              "min": 2024.0,
              "max": 2024.0,
              "std": 0.0
            },
            "sales": {
              "mean": 1199.1343749999999,
              "min": 650.5,
              "max": 2250.8,
              "std": 404.6938011030187
            },
            "orders": {
              "mean": 531.875,
              "min": 380.0,
              "max": 860.0,
              "std": 127.29030827460966
            },
            "customers": {
              "mean": 410.3125,
              "min": 290.0,
              "max": 680.0,
              "std": 100.33763565598034
            },
            "returns": {
              "mean": 24.8125,
              "min": 12.0,
              "max": 45.0,
              "std": 7.570901425001964
            }
          }
        },
        "revision_count": 3,
        "is_unknown_source": false
      }
    },
    {
      "test_id": "TQ-003-UNKNOWN",
      "query": "根据提供的表格，计算总销售额",
      "engine_type": "TableQA",
      "execution_success": true,
      "answer_correct": true,
      "confidence": 0.5,
      "execution_time_ms": 13338.815927505493,
      "error": null,
      "generated_query": "total_sales = df['sales'].sum()\nprint(total_sales)",
      "details": {
        "table_schema": {
          "columns": [
            {
              "name": "quarter",
              "dtype": "object",
              "non_null": 32,
              "unique": 4,
              "sample_values": [
                "Q1",
                "Q1",
                "Q1"
              ]
            },
            {
              "name": "year",
              "dtype": "int64",
              "non_null": 32,
              "unique": 1,
              "sample_values": [
                2024,
                2024,
                2024
              ]
            },
            {
              "name": "region",
              "dtype": "object",
              "non_null": 32,
              "unique": 4,
              "sample_values": [
                "North",
                "North",
                "South"
              ]
            },
            {
              "name": "product_category",
              "dtype": "object",
              "non_null": 32,
              "unique": 2,
              "sample_values": [
                "Electronics",
                "Clothing",
                "Electronics"
              ]
            },
            {
              "name": "sales",
              "dtype": "float64",
              "non_null": 32,
              "unique": 31,
              "sample_values": [
                1250.5,
                890.3,
                980.2
              ]
            },
            {
              "name": "orders",
              "dtype": "int64",
              "non_null": 32,
              "unique": 23,
              "sample_values": [
                450,
                520,
                380
              ]
            },
            {
              "name": "customers",
              "dtype": "int64",
              "non_null": 32,
              "unique": 20,
              "sample_values": [
                320,
                410,
                290
              ]
            },
            {
              "name": "returns",
              "dtype": "int64",
              "non_null": 32,
              "unique": 22,
              "sample_values": [
                15,
                28,
                12
              ]
            }
          ],
          "row_count": 32,
          "dtypes": {
            "quarter": "object",
            "year": "int64",
            "region": "object",
            "product_category": "object",
            "sales": "float64",
            "orders": "int64",
            "customers": "int64",
            "returns": "int64"
          },
          "statistics": {
            "year": {
              "mean": 2024.0,
              "min": 2024.0,
              "max": 2024.0,
              "std": 0.0
            },
            "sales": {
              "mean": 1199.1343749999999,
              "min": 650.5,
              "max": 2250.8,
              "std": 404.6938011030187
            },
            "orders": {
              "mean": 531.875,
              "min": 380.0,
              "max": 860.0,
              "std": 127.29030827460966
            },
            "customers": {
              "mean": 410.3125,
              "min": 290.0,
              "max": 680.0,
              "std": 100.33763565598034
            },
            "returns": {
              "mean": 24.8125,
              "min": 12.0,
              "max": 45.0,
              "std": 7.570901425001964
            }
          }
        },
        "revision_count": 3,
        "is_unknown_source": true
      }
    }
  ]
}

path: comprehensive_tests.py
content: #!/usr/bin/env python3
"""
HeteroMind Comprehensive Test Suite

Tests all three engines (SQL, SPARQL, TableQA) with accuracy metrics.
Includes ground truth comparison and unknown knowledge source testing.

Usage:
    python3 comprehensive_tests.py
"""

import asyncio
import os
import json
import time
import sys
from pathlib import Path
from typing import Dict, Any, List, Optional
from dataclasses import dataclass, field
from datetime import datetime

sys.path.insert(0, str(Path(__file__).parent / "src"))

from engines.nl2sql.multi_stage_engine import MultiStageNL2SQLEngine
from engines.nl2sparql.multi_stage_engine import MultiStageNL2SPARQLEngine
from engines.table_qa.multi_stage_engine import MultiStageTableQAEngine

DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "sk-YOUR_API_KEY_HERE")
DEEPSEEK_BASE_URL = "https://api.deepseek.com/v1"
DEEPSEEK_MODEL = "deepseek-chat"
TEST_DATA_DIR = Path(__file__).parent / "tests" / "test_data"

@dataclass
class TestResult:
    test_id: str
    query: str
    engine_type: str
    generated_query: Optional[str]
    expected_query: Optional[str]
    execution_success: bool
    answer_correct: bool
    confidence: float
    execution_time_ms: float
    error: Optional[str] = None
    details: Dict = field(default_factory=dict)

@dataclass
class TestSuiteResult:
    total_tests: int
    passed_tests: int
    failed_tests: int
    accuracy: float
    avg_confidence: float
    avg_time_ms: float
    by_engine: Dict[str, Dict]
    detailed_results: List[TestResult]

SQL_TEST_CASES = [
    {"id": "SQL-001", "query": "How many employees are in the Engineering department?", "language": "en", "expected_sql": "SELECT COUNT(*) FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'Engineering'", "expected_tables": ["employees", "departments"]},
    {"id": "SQL-002", "query": "按部门统计每个部门的平均薪资，并显示薪资最高的前 3 个部门", "language": "zh", "expected_sql": "SELECT d.name, AVG(e.salary) FROM employees e JOIN departments d ON e.department_id = d.id GROUP BY d.name ORDER BY AVG(e.salary) DESC LIMIT 3", "expected_tables": ["employees", "departments"]},
    {"id": "SQL-003", "query": "List all employees hired after 2020", "language": "en", "expected_sql": "SELECT * FROM employees WHERE hire_date > '2020-01-01'", "expected_tables": ["employees"]},
]

SPARQL_TEST_CASES = [
    {"id": "KG-001", "query": "Who is the founder of Microsoft?", "language": "en", "expected_sparql": "SELECT ?founder WHERE { <http://dbpedia.org/resource/Microsoft> <http://dbpedia.org/ontology/founder> ?founder }", "expected_entities": ["Microsoft"]},
    {"id": "KG-002", "query": "What companies are headquartered in Seattle?", "language": "en", "expected_sparql": "SELECT ?company WHERE { ?company <http://dbpedia.org/ontology/headquarters> <http://dbpedia.org/resource/Seattle> }", "expected_entities": ["Seattle"]},
]

TABLEQA_TEST_CASES = [
    {"id": "TQ-001", "query": "2024 年哪个季度的销售额最高？", "language": "zh", "expected_code_contains": ["groupby", "max", "quarter", "sales"], "expected_columns": ["quarter", "sales"], "is_unknown_source": False},
    {"id": "TQ-002", "query": "What is the average sales per quarter?", "language": "en", "expected_code_contains": ["mean", "average", "sales"], "expected_columns": ["quarter", "sales"], "is_unknown_source": False},
    {"id": "TQ-003-UNKNOWN", "query": "根据提供的表格，计算总销售额", "language": "zh", "expected_code_contains": ["sum", "sales"], "expected_columns": ["sales"], "is_unknown_source": True},
]

class HeteroMindTestSuite:
    def __init__(self):
        self.results = []
        self.sql_engine = None
        self.sparql_engine = None
        self.table_engine = None
    
    def setup_engines(self):
        print("\n" + "="*70)
        print("Initializing HeteroMind Engines")
        print("="*70)
        
        try:
            schema_path = TEST_DATA_DIR / "company_db_schema.json"
            with open(schema_path) as f:
                schema = json.load(f)
            self.sql_engine = MultiStageNL2SQLEngine({"name": "test_sql", "schema": schema, "llm_config": {"model": DEEPSEEK_MODEL, "api_key": DEEPSEEK_API_KEY, "base_url": DEEPSEEK_BASE_URL}, "generation_config": {"num_candidates": 3, "max_revisions": 2, "voting_enabled": True, "parallel_generation": True}})
            print("✅ SQL Engine initialized")
        except Exception as e:
            print(f"❌ SQL Engine failed: {e}")
        
        try:
            ontology_path = TEST_DATA_DIR / "kg_ontology.json"
            with open(ontology_path) as f:
                ontology = json.load(f)
            self.sparql_engine = MultiStageNL2SPARQLEngine({"name": "test_sparql", "endpoint_url": "https://dbpedia.org/sparql", "ontology": ontology, "llm_config": {"model": DEEPSEEK_MODEL, "api_key": DEEPSEEK_API_KEY, "base_url": DEEPSEEK_BASE_URL}, "generation_config": {"num_candidates": 3, "max_revisions": 2, "voting_enabled": True}})
            print("✅ SPARQL Engine initialized")
        except Exception as e:
            print(f"❌ SPARQL Engine failed: {e}")
        
        try:
            self.table_engine = MultiStageTableQAEngine({"name": "test_table", "table_path": str(TEST_DATA_DIR / "sales_quarterly.csv"), "llm_config": {"model": DEEPSEEK_MODEL, "api_key": DEEPSEEK_API_KEY, "base_url": DEEPSEEK_BASE_URL}, "generation_config": {"num_candidates": 3, "max_revisions": 2, "voting_enabled": True}})
            print("✅ TableQA Engine initialized")
        except Exception as e:
            print(f"❌ TableQA Engine failed: {e}")
        
        print("="*70)
    
    async def test_sql_engine(self) -> List[TestResult]:
        if not self.sql_engine:
            print("⚠️  SQL engine not available, skipping SQL tests")
            return []
        
        results = []
        print("\n" + "="*70)
        print("Testing SQL Engine")
        print("="*70)
        
        for tc in SQL_TEST_CASES:
            print(f"\n[{tc['id']}] {tc['query']}")
            start_time = time.time()
            
            try:
                result = await self.sql_engine.execute(tc["query"], {})
                exec_time = (time.time() - start_time) * 1000
                sql_correct = self._evaluate_sql_correctness(result.generated_sql, tc["expected_sql"], tc["expected_tables"])
                
                test_result = TestResult(test_id=tc["id"], query=tc["query"], engine_type="SQL", generated_query=result.generated_sql, expected_query=tc["expected_sql"], execution_success=result.success, answer_correct=sql_correct, confidence=result.final_confidence, execution_time_ms=exec_time, details={"schema_links": str(result.schema_links) if hasattr(result, "schema_links") else None, "revision_count": result.revision_count})
                
                status = "✅ PASS" if sql_correct else "❌ FAIL"
                print(f"  {status} | Confidence: {result.final_confidence:.2f} | Time: {exec_time:.0f}ms")
                print(f"  Generated: {result.generated_sql[:100]}...")
                
            except Exception as e:
                test_result = TestResult(test_id=tc["id"], query=tc["query"], engine_type="SQL", generated_query=None, expected_query=tc["expected_sql"], execution_success=False, answer_correct=False, confidence=0.0, execution_time_ms=(time.time() - start_time) * 1000, error=str(e))
                print(f"  ❌ ERROR: {e}")
            
            results.append(test_result)
        
        return results
    
    async def test_sparql_engine(self) -> List[TestResult]:
        if not self.sparql_engine:
            print("⚠️  SPARQL engine not available, skipping SPARQL tests")
            return []
        
        results = []
        print("\n" + "="*70)
        print("Testing SPARQL Engine")
        print("="*70)
        
        for tc in SPARQL_TEST_CASES:
            print(f"\n[{tc['id']}] {tc['query']}")
            start_time = time.time()
            
            try:
                result = await self.sparql_engine.execute(tc["query"], {})
                exec_time = (time.time() - start_time) * 1000
                sparql_correct = self._evaluate_sparql_correctness(result.generated_sparql, tc["expected_sparql"], tc["expected_entities"])
                
                test_result = TestResult(test_id=tc["id"], query=tc["query"], engine_type="SPARQL", generated_query=result.generated_sparql, expected_query=tc["expected_sparql"], execution_success=result.success, answer_correct=sparql_correct, confidence=result.final_confidence, execution_time_ms=exec_time, details={"linked_entities": str(result.linked_entities) if hasattr(result, "linked_entities") else None, "revision_count": result.revision_count})
                
                status = "✅ PASS" if sparql_correct else "❌ FAIL"
                print(f"  {status} | Confidence: {result.final_confidence:.2f} | Time: {exec_time:.0f}ms")
                print(f"  Generated: {result.generated_sparql[:100]}...")
                
            except Exception as e:
                test_result = TestResult(test_id=tc["id"], query=tc["query"], engine_type="SPARQL", generated_query=None, expected_query=tc["expected_sparql"], execution_success=False, answer_correct=False, confidence=0.0, execution_time_ms=(time.time() - start_time) * 1000, error=str(e))
                print(f"  ❌ ERROR: {e}")
            
            results.append(test_result)
        
        return results
    
    async def test_table_engine(self) -> List[TestResult]:
        if not self.table_engine:
            print("⚠️  TableQA engine not available, skipping TableQA tests")
            return []
        
        results = []
        print("\n" + "="*70)
        print("Testing TableQA Engine")
        print("="*70)
        
        for tc in TABLEQA_TEST_CASES:
            unknown_tag = " [UNKNOWN SOURCE]" if tc.get("is_unknown_source") else ""
            print(f"\n[{tc['id']}]{unknown_tag} {tc['query']}")
            start_time = time.time()
            
            try:
                result = await self.table_engine.execute(tc["query"], {})
                exec_time = (time.time() - start_time) * 1000
                code_correct = self._evaluate_code_correctness(result.generated_code, tc["expected_code_contains"], tc["expected_columns"])
                
                test_result = TestResult(test_id=tc["id"], query=tc["query"], engine_type="TableQA", generated_query=result.generated_code, expected_query=tc.get("expected_code_contains"), execution_success=result.success, answer_correct=code_correct, confidence=result.final_confidence, execution_time_ms=exec_time, details={"table_schema": dict(result.table_schema) if isinstance(result.table_schema, dict) else str(result.table_schema) if result.table_schema else None, "revision_count": result.revision_count, "is_unknown_source": tc.get("is_unknown_source", False)})
                
                status = "✅ PASS" if code_correct else "❌ FAIL"
                unknown_info = " (Unknown Source)" if tc.get("is_unknown_source") else ""
                print(f"  {status}{unknown_info} | Confidence: {result.final_confidence:.2f} | Time: {exec_time:.0f}ms")
                print(f"  Generated: {result.generated_code[:100]}...")
                
            except Exception as e:
                test_result = TestResult(test_id=tc["id"], query=tc["query"], engine_type="TableQA", generated_query=None, expected_query=tc.get("expected_code_contains"), execution_success=False, answer_correct=False, confidence=0.0, execution_time_ms=(time.time() - start_time) * 1000, error=str(e))
                print(f"  ❌ ERROR: {e}")
            
            results.append(test_result)
        
        return results
    
    def _evaluate_sql_correctness(self, generated: str, expected: str, expected_tables: List[str]) -> bool:
        if not generated:
            return False
        has_tables = all(table.lower() in generated.lower() for table in expected_tables)
        has_select = "SELECT" in generated.upper()
        return has_tables and has_select
    
    def _evaluate_sparql_correctness(self, generated: str, expected: str, expected_entities: List[str]) -> bool:
        if not generated:
            return False
        has_entities = any(entity.lower() in generated.lower() for entity in expected_entities)
        has_select = "SELECT" in generated.upper()
        has_where = "WHERE" in generated.upper()
        return has_entities and has_select and has_where
    
    def _evaluate_code_correctness(self, generated: str, expected_contains: List[str], expected_columns: List[str]) -> bool:
        if not generated:
            return False
        has_patterns = any(pattern.lower() in generated.lower() for pattern in expected_contains)
        has_columns = any(col.lower() in generated.lower() for col in expected_columns)
        return has_patterns and has_columns
    
    def calculate_summary(self, results: List[TestResult]) -> TestSuiteResult:
        total = len(results)
        passed = sum(1 for r in results if r.answer_correct)
        failed = total - passed
        
        by_engine = {}
        for engine_type in ["SQL", "SPARQL", "TableQA"]:
            engine_results = [r for r in results if r.engine_type == engine_type]
            if engine_results:
                engine_passed = sum(1 for r in engine_results if r.answer_correct)
                by_engine[engine_type] = {"total": len(engine_results), "passed": engine_passed, "accuracy": engine_passed / len(engine_results), "avg_confidence": sum(r.confidence for r in engine_results) / len(engine_results), "avg_time_ms": sum(r.execution_time_ms for r in engine_results) / len(engine_results)}
        
        return TestSuiteResult(total_tests=total, passed_tests=passed, failed_tests=failed, accuracy=passed / total if total > 0 else 0, avg_confidence=sum(r.confidence for r in results) / total if total > 0 else 0, avg_time_ms=sum(r.execution_time_ms for r in results) / total if total > 0 else 0, by_engine=by_engine, detailed_results=results)
    
    def print_summary(self, summary: TestSuiteResult):
        print("\n" + "="*70)
        print("TEST SUMMARY")
        print("="*70)
        print(f"Total Tests:     {summary.total_tests}")
        print(f"Passed:          {summary.passed_tests} ✅")
        print(f"Failed:          {summary.failed_tests} ❌")
        print(f"Overall Accuracy: {summary.accuracy:.1%}")
        print(f"Avg Confidence:  {summary.avg_confidence:.2f}")
        print(f"Avg Time:        {summary.avg_time_ms:.0f}ms")
        
        print("\nBy Engine:")
        for engine, stats in summary.by_engine.items():
            print(f"\n  {engine}:")
            print(f"    Tests:    {stats['total']}")
            print(f"    Passed:   {stats['passed']}/{stats['total']} ({stats['accuracy']:.1%})")
            print(f"    Confidence: {stats['avg_confidence']:.2f}")
            print(f"    Avg Time: {stats['avg_time_ms']:.0f}ms")
        
        print("\n" + "="*70)
    
    def save_results(self, summary: TestSuiteResult, output_path: str):
        output = {"timestamp": datetime.now().isoformat(), "summary": {"total_tests": summary.total_tests, "passed_tests": summary.passed_tests, "failed_tests": summary.failed_tests, "accuracy": summary.accuracy, "avg_confidence": summary.avg_confidence, "avg_time_ms": summary.avg_time_ms}, "by_engine": summary.by_engine, "detailed_results": [{"test_id": r.test_id, "query": r.query, "engine_type": r.engine_type, "execution_success": r.execution_success, "answer_correct": r.answer_correct, "confidence": r.confidence, "execution_time_ms": r.execution_time_ms, "error": r.error, "generated_query": r.generated_query, "details": r.details} for r in summary.detailed_results]}
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(output, f, indent=2, ensure_ascii=False)
        
        print(f"\n📄 Results saved to: {output_path}")

async def main():
    print("\n" + "="*70)
    print("HeteroMind Comprehensive Test Suite")
    print("="*70)
    print(f"Model: {DEEPSEEK_MODEL}")
    print(f"Test Cases: {len(SQL_TEST_CASES)} SQL + {len(SPARQL_TEST_CASES)} SPARQL + {len(TABLEQA_TEST_CASES)} TableQA")
    print("="*70)
    
    suite = HeteroMindTestSuite()
    suite.setup_engines()
    
    all_results = []
    all_results.extend(await suite.test_sql_engine())
    all_results.extend(await suite.test_sparql_engine())
    all_results.extend(await suite.test_table_engine())
    
    summary = suite.calculate_summary(all_results)
    suite.print_summary(summary)
    
    output_path = Path(__file__).parent / "tests" / "comprehensive_test_results.json"
    suite.save_results(summary, str(output_path))
    
    report_path = Path(__file__).parent / "tests" / "comprehensive_test_report.md"
    generate_markdown_report(summary, str(report_path))
    
    return summary

def generate_markdown_report(summary: TestSuiteResult, output_path: str):
    report = f"""# HeteroMind Comprehensive Test Report

**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
**Model:** {DEEPSEEK_MODEL}

## Summary

| Metric | Value |
|--------|-------|
| Total Tests | {summary.total_tests} |
| Passed | {summary.passed_tests} ✅ |
| Failed | {summary.failed_tests} ❌ |
| **Overall Accuracy** | **{summary.accuracy:.1%}** |
| Avg Confidence | {summary.avg_confidence:.2f} |
| Avg Response Time | {summary.avg_time_ms:.0f}ms |

## Results by Engine

| Engine | Tests | Passed | Accuracy | Avg Confidence | Avg Time |
|--------|-------|--------|----------|----------------|----------|
"""
    for engine, stats in summary.by_engine.items():
        report += f"| {engine} | {stats['total']} | {stats['passed']} | {stats['accuracy']:.1%} | {stats['avg_confidence']:.2f} | {stats['avg_time_ms']:.0f}ms |\n"
    
    report += "\n## Detailed Results\n\n"
    
    for result in summary.detailed_results:
        status = "✅ PASS" if result.answer_correct else "❌ FAIL"
        unknown_tag = " [UNKNOWN SOURCE]" if result.details.get("is_unknown_source") else ""
        report += f"### {result.test_id}{unknown_tag} - {status}\n\n**Query:** {result.query}\n\n**Engine:** {result.engine_type}\n\n**Confidence:** {result.confidence:.2f}\n\n**Execution Time:** {result.execution_time_ms:.0f}ms\n\n**Generated Query:**\n```\n{result.generated_query or 'N/A'}\n```\n\n"
        if result.error:
            report += f"**Error:** {result.error}\n\n"
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(report)
    
    print(f"📄 Report saved to: {output_path}")

if __name__ == "__main__":
    asyncio.run(main())


path: README.md
content: # HeteroMind

**Unified Heterogeneous Knowledge QA System**

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![ClawHub](https://img.shields.io/badge/ClawHub-skill-green.svg)](https://clawhub.ai/)

HeteroMind is a multi-source knowledge QA system that automatically routes natural language queries to appropriate knowledge sources (SQL databases, Knowledge Graphs, Table files) without requiring users to specify the data source.

---

## ✨ Features

| Feature | Description |
|---------|-------------|
| **Auto Source Detection** | 4-layer detection architecture automatically identifies SQL, SPARQL, or Table QA |
| **Multi-language Support** | Supports both Chinese and English queries |
| **Multi-stage Reasoning** | Schema Linking → Query Generation → Self-Revision → Voting → Execution |
| **Parallel Execution** | Concurrent candidate generation for improved efficiency |
| **Self-Revision** | Multi-round review and improvement of generated queries |
| **Voting Mechanism** | Multi-candidate voting to select the best result |

---

## 🏗️ Architecture

### 4-Layer Source Detection

```
User Query
    ↓
┌─────────────────────────────────────────┐
│ Layer 1: Rule-Based (15%)                │ ← Keywords, patterns
├─────────────────────────────────────────┤
│ Layer 2: LLM Semantic (35%)              │ ← Intent understanding
├─────────────────────────────────────────┤
│ Layer 3a: SQL Schema Match (25%)         │ ← Database schema matching
├─────────────────────────────────────────┤
│ Layer 3b: KG Entity Link (25%)           │ ← Entity linking to KG
├─────────────────────────────────────────┤
│ Layer 3c: Entity Verification (25%+30%)  │ ← Verify entities exist
└─────────────────────────────────────────┘
    ↓
Layer 4: Multi-Source Fusion → Final Decision + Execution Plan
```

### Query Execution Pipeline

```
Source Selected
    ↓
1. Schema/Entity Linking
2. Parallel Candidate Generation (3 candidates)
3. Multi-Round Self-Revision (2 rounds)
4. Validation (syntax & semantic)
5. Voting (select best candidate)
6. Execution (run query)
7. Result Verification
    ↓
Natural Language Answer
```

---

## 📦 Installation

```bash
# Clone or download the skill
cd HeteroMind

# Install dependencies
pip install -r requirements.txt
```

### Requirements

- Python 3.10+
- `aiohttp` - Async HTTP client
- `pandas` - Data manipulation
- `openpyxl` - Excel support
- `openai` - LLM API client (optional, for LLM-based detection)

---

## 🚀 Quick Start

### Basic Usage

```python
import asyncio
from src.orchestrator import HeteroMindOrchestrator

# Configuration
config = {
    "source_detection": {
        "layer2": {
            "api_key": "your-api-key",
            "model": "deepseek-chat",
            "base_url": "https://api.deepseek.com/v1",
        },
        "layer3": {
            "schemas": [{"database": "mydb", "tables": [...]}],
            "kg_endpoints": [
                {"name": "dbpedia", "url": "https://dbpedia.org/sparql"}
            ],
        },
    },
    "engines": {
        "sql": [{"name": "default", "enabled": True}],
        "sparql": [{"name": "default", "enabled": True}],
        "table_qa": [{"name": "default", "enabled": True}],
    },
}

async def main():
    orchestrator = HeteroMindOrchestrator(config)
    
    # Execute query with default LLM
    response = await orchestrator.query("How many employees are in Engineering?")
    print(f"Answer: {response.answer}")
    print(f"Source: {response.sources}")
    print(f"Confidence: {response.confidence:.2f}")
    
    # Execute query with specific model override
    response = await orchestrator.query(
        query="Show total sales by region",
        model="gpt-4-turbo",      # Override model
        [REDACTED_SECRET]",         # Override API key
    )

asyncio.run(main())
```

### Source Detection Only

```python
from src.classifier import SourceDetectorOrchestrator
import json

# Load schema
with open('tests/test_data/company_db_schema.json') as f:
    schema = json.load(f)

# Create detector
detector = SourceDetectorOrchestrator({
    "layer2": {
        "api_key": "sk-...",
        "model": "deepseek-chat",
    },
    "layer3": {
        "schemas": [schema],
    },
})

# Detect source
decision = await detector.detect("How many employees in Engineering?")

print(f"Primary Source: {decision.primary_source.value}")
print(f"Confidence: {decision.confidence:.2f}")
print(f"Reasoning: {decision.reasoning}")
```

---

## 📊 Supported Knowledge Sources

### 1. SQL Database (NL2SQL)

Supports MySQL, PostgreSQL, SQLite, and other relational databases.

**Example Queries:**
- "How many employees are in the Engineering department?"
- "Show average salary by department, top 3"

**Multi-stage Flow:**
1. Schema Linking - Identify relevant tables/columns
2. Query Generation - Generate 3 candidate SQL queries in parallel
3. Self-Revision - Multi-round review and improvement
4. Voting - Select the best candidate
5. Execution - Run the query

### 2. Knowledge Graph (NL2SPARQL)

Supports any SPARQL endpoint: DBpedia, Wikidata, custom KGs.

**Example Queries:**
- "Who is the founder of Microsoft?"
- "What companies are headquartered in Seattle?"

**Multi-stage Flow:**
1. Entity Linking - Link entities to KG
2. Ontology Retrieval - Get relevant predicates
3. Query Generation - Generate candidate SPARQL
4. Self-Revision - Review and improve
5. Voting - Select the best candidate

### 3. Table Files (TableQA)

Supports CSV, Excel, and other tabular files.

**Example Queries:**
- "Which quarter had the highest sales in 2024?"
- "Show total sales by region"

**Multi-stage Flow:**
1. Schema Analysis - Analyze table structure
2. Query Interpretation - Understand query intent
3. Code Generation - Generate Pandas code
4. Self-Revision - Review and improve code
5. Execution - Safely execute code

---

## 🧪 Testing

### Run Test Suite

```bash
cd HeteroMind
python3 comprehensive_tests.py
```

This will:
- Test all 3 engines (SQL, SPARQL, TableQA)
- Run 8 test cases total
- Generate detailed reports:
  - `tests/comprehensive_test_results.json` - JSON format results
  - `tests/comprehensive_test_report.md` - Markdown format report

### Test Data

Test data is located in `tests/test_data/`:

| File | Description |
|------|-------------|
| `company_db_schema.json` | SQL database schema |
| `sales_quarterly.csv` | Sample sales data (32 rows) |
| `kg_ontology.json` | Knowledge graph ontology |

---

## 🔧 Configuration

### LLM Configuration

Supports multiple LLM providers. Configure at initialization or per-call:

#### Global Configuration (Engine Initialization)

```python
from src.engines.nl2sql.multi_stage_engine import MultiStageNL2SQLEngine

# DeepSeek
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "deepseek-chat",
        "api_key": "sk-...",
        "base_url": "https://api.deepseek.com/v1",
    },
})

# GPT-4
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "gpt-4",
        "api_key": "sk-...",
    },
})

# Claude (via Anthropic)
engine = MultiStageNL2SQLEngine({
    "name": "sql_engine",
    "schema": schema,
    "llm_config": {
        "model": "claude-3-sonnet-20240229",
        "api_key": "sk-ant-...",
    },
})
```

#### Per-Call Override (Runtime)

```python
# Override model and API key at execution time
result = await engine.execute(
    query="How many employees?",
    context={},
    model="gpt-4-turbo",      # Override model
    [REDACTED_SECRET]",         # Override API key
)
```

#### Supported Providers

| Provider | Models | Configuration |
|----------|--------|---------------|
| **DeepSeek** | deepseek-chat, deepseek-coder | `base_url: https://api.deepseek.com/v1` |
| **OpenAI** | gpt-4, gpt-3.5-turbo | Default endpoint |
| **Azure OpenAI** | gpt-4, gpt-35-turbo | `base_url: https://{resource}.openai.azure.com` |
| **Anthropic** | claude-3-opus, claude-3-sonnet | Use OpenAI-compatible wrapper |
| **Local LLM** | Any OpenAI-compatible | `base_url: http://localhost:11434/v1` (Ollama) |


### 🔒 Security Best Practices

**Never commit API keys to version control!**

```bash
# ✅ Recommended: Use environment variables
export DEEPSEEK_API_KEY="sk-..."
export OPENAI_API_KEY="sk-..."

# ✅ Use .env file (add to .gitignore)
echo "DEEPSEEK_API_KEY=sk-..." >> .env

# ❌ Never hardcode API keys in code
llm_config = {"api_key": "sk-REAL-KEY-HERE"}  # DON'T DO THIS
```

**API Key Protection:**
- HeteroMind automatically sanitizes outputs to prevent credential leakage
- LLM prompts include security instructions to never output API keys
- Use environment variables or secure vaults for credential storage

### Generation Parameters

```python
generation_config = {
    "num_candidates": 3,        # Number of candidates to generate
    "max_revisions": 2,         # Maximum revision rounds
    "parallel_generation": True, # Enable parallel generation
    "voting_enabled": True,     # Enable voting mechanism
}
```

### Source Detection Weights

```python
weights = {
    "rule_based": 0.15,      # Layer 1: Keyword/pattern matching
    "llm_based": 0.35,       # Layer 2: LLM semantic classification
    "schema_based": 0.25,    # Layer 3a/3b: Schema/entity matching
    "verification": 0.25,    # Layer 3c: Entity verification
}
```

---

## 📁 Project Structure

```
HeteroMind/
├── README.md                    # This file
├── SKILL.md                     # ClawHub skill descriptor
├── requirements.txt             # Python dependencies
├── comprehensive_tests.py       # Comprehensive test suite
├── config/
│   ├── source_detection.yaml    # Source detection configuration
│   └── sources.yaml             # Knowledge source connections
├── src/
│   ├── __init__.py
│   ├── orchestrator.py          # Main orchestrator
│   ├── classifier/              # 4-layer source detection
│   │   ├── models.py            # Data models
│   │   ├── rule_detector.py     # Layer 1: Rule-based
│   │   ├── llm_detector.py      # Layer 2: LLM-based
│   │   ├── sql_schema_matcher.py # Layer 3a: SQL schema
│   │   ├── kg_entity_linker.py  # Layer 3b: KG entity linking
│   │   ├── entity_verifier.py   # Layer 3c: Entity verification
│   │   ├── source_fusion.py     # Layer 4: Multi-source fusion
│   │   └── orchestrator.py      # Detection orchestrator
│   ├── decomposer/              # Task decomposition
│   ├── engines/                 # Query engines
│   │   ├── base.py              # Base classes
│   │   ├── nl2sql/              # NL2SQL engine
│   │   ├── nl2sparql/           # NL2SPARQL engine
│   │   └── table_qa/            # TableQA engine
│   ├── fusion/                  # Result fusion
│   └── generator/               # Answer generation
└── tests/
    ├── README.md                # Test documentation
    └── test_data/               # Test data files
```

---

## 📈 Performance

| Metric | Value |
|--------|-------|
| Average Response Time (SQL) | ~22s |
| Average Response Time (KG/Table) | ~30s |
| Source Detection Time | <0.1s |
| Candidates Generated | 3 per query |
| Revision Rounds | 2 (configurable) |
| Confidence Range | 0.2 - 1.0 |

### Test Results

| Engine | Tests | Passed | Accuracy | Avg Confidence |
|--------|-------|--------|----------|----------------|
| **SQL (NL2SQL)** | 3 | 3 | 100.0% | 0.60 |
| **SPARQL (NL2SPARQL)** | 2 | 2 | 100.0% | 0.20 |
| **TableQA** | 3 | 3 | 100.0% | 0.62 |
| **Overall** | 8 | 8 | 100.0% | 0.51 |

---

## 👥 Author

**Coin Lab, Southeast University**

---

## 🙏 Acknowledgments

- DeepSeek for providing LLM API
- DBpedia and Wikidata for open knowledge graphs
- OpenClaw community for skill development framework



path: SECURITY.md
content: # Security & Safety Documentation

## Overview

HeteroMind is a powerful knowledge QA system that can:
- Generate and execute SQL queries against databases
- Generate and execute SPARQL queries against knowledge graphs
- Generate and execute Python (pandas) code for table analysis
- Scan local file systems for table data
- Connect to external LLM APIs and SPARQL endpoints

This document describes the security model, risks, and mitigation measures.

---

## Threat Model

### Assets at Risk

| Asset | Risk | Impact |
|-------|------|--------|
| **Database Credentials** | Theft via logs/output | Unauthorized data access |
| **API Keys** | Leakage in generated code | Account compromise, billing fraud |
| **Local Files** | Unauthorized read access | Data exfiltration |
| **Network Endpoints** | Unrestricted access | Data leakage, resource abuse |

### Attack Vectors

1. **Prompt Injection** - Malicious input causing unintended queries
2. **Credential Leakage** - API keys exposed in output or logs
3. **SQL Injection** - Generated SQL with injection vulnerabilities
4. **Path Traversal** - File scanning accessing unintended directories
5. **Code Injection** - Generated Python code with malicious payloads

---

## Security Controls

### 1. Credential Protection

**Implementation:** `src/utils/api_security.py`

```python
# All outputs are sanitized
from src.utils.api_security import sanitize_output

safe_output = sanitize_output(llm_response)

# API keys validated before use
from src.utils.api_security import validate_no_api_key_in_response

if validate_no_api_key_in_response(generated_code):
    execute(generated_code)
```

**Best Practices:**
- Never hardcode credentials in code or config files
- Use environment variables exclusively
- Rotate API keys regularly
- Use read-only database credentials when possible

### 2. Sandboxed Execution

**Python Code Execution:**
```python
# Restricted scope - no arbitrary imports
safe_scope = {
    "df": dataframe,
    "pd": pandas,  # Only pandas, no os/sys/subprocess
    "np": numpy,
}

# Execute in restricted scope
exec(generated_code, safe_scope)
```

**SQL Execution:**
- Read-only queries only (SELECT)
- DDL/DML operations (DROP, DELETE, UPDATE) blocked
- Query timeout enforced (default: 30s)

### 3. User Consent

**Auto-Execution Control:**

```yaml
# config/source_detection.yaml
auto_execute: false         # Default: disabled
require_confirmation: true  # Default: enabled

thresholds:
  high_confidence: 0.95     # Very conservative
```

**Confirmation Flow:**
```
Query: "Show all employees"
Detection Confidence: 0.87

⚠️  Confidence below threshold (0.95)
Generated SQL: SELECT * FROM employees

Proceed with execution? [y/N]
```

### 4. File Access Restrictions

**Default Configuration:**
```yaml
entity_verification:
  tables:
    enabled: false          # Disabled by default
    paths: []               # No automatic paths
    explicit_paths_only: true
```

**Explicit Path Configuration:**
```yaml
entity_verification:
  tables:
    enabled: true
    paths:
      - "/home/user/data/sales_2024.csv"  # Explicit file, no wildcards
    max_file_size_mb: 10
```

### 5. Audit Logging

**Log Configuration:**
```yaml
logging:
  level: "INFO"
  log_layer_outputs: true
  log_verification_details: true
  log_generated_queries: true  # Log all generated SQL/SPARQL/code
  log_execution_results: true  # Log execution outcomes
```

**Log Output:**
```json
{
  "timestamp": "2026-04-12T10:30:00Z",
  "query": "How many employees?",
  "detected_source": "sql_database",
  "confidence": 0.92,
  "generated_sql": "SELECT COUNT(*) FROM employees",
  "execution_result": "success",
  "row_count": 150,
  "execution_time_ms": 234
}
```

---

## Configuration Profiles

### Profile 1: Maximum Security (Recommended for Production)

```yaml
# source_detection.yaml
auto_execute: false
require_confirmation: true

thresholds:
  high_confidence: 0.98
  medium_confidence: 0.8
  low_confidence: 0.6

entity_verification:
  tables:
    enabled: false
  sql:
    enabled: true
    read_only: true
  kg:
    enabled: true
    use_ask_query: true

execution:
  timeout:
    total_seconds: 30
    per_source_seconds: 10
```

### Profile 2: Balanced (Development/Testing)

```yaml
auto_execute: false
require_confirmation: true

thresholds:
  high_confidence: 0.9
  medium_confidence: 0.7
  low_confidence: 0.5

entity_verification:
  tables:
    enabled: true
    paths:
      - "${WORKSPACE}/test_data/*.csv"
    max_file_size_mb: 50
```

### Profile 3: Permissive (Local Development Only)

```yaml
auto_execute: true
require_confirmation: false

thresholds:
  high_confidence: 0.8
  medium_confidence: 0.5
  low_confidence: 0.3

# ⚠️  Only use in isolated development environments!
```

---

## Environment Variables

### Required

| Variable | Purpose | Security Level |
|----------|---------|----------------|
| `DEEPSEEK_API_KEY` | LLM query generation | High - protect like password |
| `OPENAI_API_KEY` | Alternative LLM | High - protect like password |

### Optional

| Variable | Purpose | Security Level |
|----------|---------|----------------|
| `MYSQL_CONNECTION_STRING` | Database access | Critical - use read-only creds |
| `CUSTOM_KG_ENDPOINT` | KG endpoint | Medium - verify endpoint trust |
| `WORKSPACE` | File scanning base | Medium - restrict to specific dirs |

### Setup

```bash
# Create .env file (NEVER commit to git)
cp .env.example .env
chmod 600 .env  # Restrict permissions

# Edit with your credentials
nano .env

# Load environment
export $(cat .env | xargs)

# Verify
echo $DEEPSEEK_API_KEY  # Should show value
git status .env         # Should show as untracked
```

---

## Incident Response

### If Credentials Are Compromised

1. **Immediately rotate** the affected API key or password
2. **Review audit logs** for unauthorized access
3. **Revoke access** for any suspicious sessions
4. **Update .env** file with new credentials
5. **Notify affected parties** if user data was accessed

### If Malicious Query Detected

1. **Stop execution** immediately (Ctrl+C or kill process)
2. **Review generated query** in audit logs
3. **Identify injection vector** in user input
4. **Add input validation** to prevent recurrence
5. **Report security issue** to maintainers

---

## Compliance Considerations

### Data Protection

- **GDPR**: Ensure personal data in databases is handled per GDPR requirements
- **CCPA**: Provide opt-out mechanisms for California residents
- **HIPAA**: Do not use with PHI unless properly configured and audited

### Access Control

- Use role-based database credentials (read-only for QA)
- Implement API key rotation policies
- Enable audit logging for compliance reporting

---

## Security Checklist

Before deploying HeteroMind:

- [ ] All credentials stored in environment variables (not code)
- [ ] `.env` file added to `.gitignore`
- [ ] `auto_execute: false` in production config
- [ ] `require_confirmation: true` in production config
- [ ] Database credentials are read-only
- [ ] File scanning restricted to specific directories
- [ ] Audit logging enabled
- [ ] API keys rotated within last 90 days
- [ ] Network endpoints whitelisted in firewall
- [ ] Security review completed

---

## Reporting Security Issues

If you discover a security vulnerability:

1. **Do not** create a public GitHub issue
2. **Email**: [REDACTED_SECRET] (if applicable)
3. **Include**: Description, reproduction steps, impact assessment
4. **Wait**: Allow 90 days for patch before public disclosure

---

## Version History

| Version | Date | Changes |
|---------|------|---------|
| 0.1.0 | 2026-04-12 | Initial security documentation |

---

*Last Updated: 2026-04-12*  
*Version: 0.1.0*


path: src/fusion/result_fusion.py
content: """
HeteroMind - Result Fusion

Fuses results from multiple knowledge sources.
"""

import logging
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from ..classifier import FinalDecision, KnowledgeSource
from ..engines.base import QueryResult

logger = logging.getLogger(__name__)


@dataclass
class FusedResult:
    """Fused result from multiple sources"""
    data: Any
    success: bool
    sources: List[str]
    confidence: float
    conflicts: List[Dict]
    fusion_method: str


class ResultFusion:
    """
    Result fusion for multi-source queries.
    
    Combines results from multiple knowledge sources,
    handles conflicts, and produces unified output.
    """
    
    def __init__(self, config: dict):
        """
        Initialize fusion module.
        
        Args:
            config: Configuration dict
        """
        self.config = config
        self.conflict_strategy = config.get("conflict_strategy", "weighted")
    
    async def fuse(
        self,
        results: List[QueryResult],
        source_decision: FinalDecision,
    ) -> FusedResult:
        """
        Fuse results from multiple sources.
        
        Args:
            results: List of QueryResult from engines
            source_decision: Source detection decision
            
        Returns:
            FusedResult
        """
        logger.info(f"Fusing {len(results)} results")
        
        if not results:
            return FusedResult(
                data=None,
                success=False,
                sources=[],
                confidence=0.0,
                conflicts=[],
                fusion_method="none",
            )
        
        # Single source: no fusion needed
        if len(results) == 1:
            return FusedResult(
                data=results[0].data,
                success=results[0].success,
                sources=[results[0].engine],
                confidence=results[0].confidence,
                conflicts=[],
                fusion_method="single",
            )
        
        # Multiple sources: fuse results
        fused_data = self._combine_data(results)
        conflicts = self._detect_conflicts(results)
        confidence = self._calculate_confidence(results)
        
        sources = list(set(r.engine for r in results))
        
        return FusedResult(
            data=fused_data,
            success=True,
            sources=sources,
            confidence=confidence,
            conflicts=conflicts,
            fusion_method=self.conflict_strategy,
        )
    
    def _combine_data(self, results: List[QueryResult]) -> Any:
        """Combine data from multiple results"""
        # Simple concatenation for now
        # In real implementation, implement smart fusion logic
        
        combined = []
        for result in results:
            if result.data:
                if isinstance(result.data, list):
                    combined.extend(result.data)
                else:
                    combined.append(result.data)
        
        return combined
    
    def _detect_conflicts(self, results: List[QueryResult]) -> List[Dict]:
        """Detect conflicts between results"""
        conflicts = []
        
        # Simple conflict detection
        # In real implementation, implement sophisticated conflict detection
        
        return conflicts
    
    def _calculate_confidence(self, results: List[QueryResult]) -> float:
        """Calculate combined confidence score"""
        if not results:
            return 0.0
        
        # Average confidence
        total = sum(r.confidence for r in results)
        return total / len(results)


path: src/fusion/__init__.py
content: """
HeteroMind - Fusion Package
"""

from .result_fusion import ResultFusion, FusedResult

__all__ = [
    "ResultFusion",
    "FusedResult",
]


path: src/decomposer/__init__.py
content: """
HeteroMind - Decomposer Package
"""

from .task_decomposer import TaskDecomposer, SubTask, TaskType

__all__ = [
    "TaskDecomposer",
    "SubTask",
    "TaskType",
]


path: src/decomposer/task_decomposer.py
content: """
HeteroMind - Task Decomposer

Breaks complex queries into executable sub-tasks.
"""

import logging
from dataclasses import dataclass
from typing import Dict, List, Optional
from enum import Enum

from ..classifier import FinalDecision, KnowledgeSource

logger = logging.getLogger(__name__)


class TaskType(Enum):
    """Sub-task types"""
    QUERY = "query"           # Simple query
    AGGREGATE = "aggregate"    # Aggregation
    JOIN = "join"             # Join operation
    FILTER = "filter"         # Filter operation
    TRANSFORM = "transform"    # Transform/format
    FUSE = "fuse"             # Fuse results


@dataclass
class SubTask:
    """Represents a sub-task in the execution plan"""
    id: str
    query: str
    source: str
    task_type: TaskType
    depends_on: List[str]
    context: dict
    priority: int = 0


class TaskDecomposer:
    """
    Task decomposer for complex queries.
    
    Breaks down multi-hop queries into executable sub-tasks
    with dependency tracking for proper execution ordering.
    """
    
    def __init__(self, config: dict):
        """
        Initialize decomposer.
        
        Args:
            config: Configuration dict
        """
        self.config = config
        self.max_depth = config.get("max_depth", 5)
    
    async def decompose(
        self,
        query: str,
        source_decision: FinalDecision,
    ) -> List[SubTask]:
        """
        Decompose query into sub-tasks.
        
        Args:
            query: Natural language query
            source_decision: Source detection decision
            
        Returns:
            List of SubTask objects
        """
        logger.info(f"Decomposing query: {query}")
        
        # Simple queries: single task
        if not source_decision.layer2_output.requires_multi_hop:
            return await self._create_simple_tasks(query, source_decision)
        
        # Complex queries: multi-step decomposition
        return await self._create_complex_tasks(query, source_decision)
    
    async def _create_simple_tasks(
        self,
        query: str,
        source_decision: FinalDecision,
    ) -> List[SubTask]:
        """Create simple single-step tasks"""
        tasks = []
        
        for i, source in enumerate(source_decision.selected_sources):
            task = SubTask(
                id=f"task_{i}",
                query=query,
                source=source.value,
                task_type=TaskType.QUERY,
                depends_on=[],
                context={
                    "entities": source_decision.layer2_output.detected_entities,
                    "predicates": source_decision.layer2_output.detected_predicates,
                },
                priority=i,
            )
            tasks.append(task)
        
        return tasks
    
    async def _create_complex_tasks(
        self,
        query: str,
        source_decision: FinalDecision,
    ) -> List[SubTask]:
        """Create complex multi-step tasks"""
        # In real implementation, use LLM to decompose
        # For now, return placeholder
        
        tasks = []
        
        # Example decomposition for hybrid query
        if KnowledgeSource.HYBRID in source_decision.selected_sources:
            # Task 1: SQL query
            tasks.append(SubTask(
                id="task_sql_1",
                query=query,  # Would be refined
                source="sql_database",
                task_type=TaskType.QUERY,
                depends_on=[],
                context={},
                priority=1,
            ))
            
            # Task 2: KG query (depends on SQL results)
            tasks.append(SubTask(
                id="task_kg_2",
                query=query,  # Would be refined
                source="knowledge_graph",
                task_type=TaskType.QUERY,
                depends_on=["task_sql_1"],
                context={},
                priority=2,
            ))
            
            # Task 3: Fuse results
            tasks.append(SubTask(
                id="task_fuse_3",
                query="fuse",
                source="fusion",
                task_type=TaskType.FUSE,
                depends_on=["task_sql_1", "task_kg_2"],
                context={},
                priority=3,
            ))
        
        return tasks


path: src/__init__.py
content: """
HeteroMind - Unified Heterogeneous Knowledge QA System

A unified natural language QA system that automatically routes queries
to appropriate knowledge sources (SQL databases, knowledge graphs,
table files, documents) without requiring users to specify the source.
"""

from .orchestrator import HeteroMindOrchestrator, QueryContext, QueryResponse, ask
from .classifier import (
    SourceDetectorOrchestrator,
    KnowledgeSource,
    FinalDecision,
)
from .engines import (
    BaseEngine,
    QueryResult,
)
from .fusion import (
    ResultFusion,
    FusedResult,
)
from .generator import (
    AnswerGenerator,
)
from .decomposer import (
    TaskDecomposer,
    SubTask,
)

__version__ = "0.1.0"
__author__ = "Coin Lab, Southeast University"

__all__ = [
    # Main orchestrator
    "HeteroMindOrchestrator",
    "QueryContext",
    "QueryResponse",
    "ask",
    # Source detection
    "SourceDetectorOrchestrator",
    "KnowledgeSource",
    "FinalDecision",
    # Engines
    "BaseEngine",
    "QueryResult",
    # Fusion
    "ResultFusion",
    "FusedResult",
    # Generator
    "AnswerGenerator",
    # Decomposer
    "TaskDecomposer",
    "SubTask",
]


path: src/utils/schema_utils.py
content: """
HeteroMind - Schema Utilities

Utilities for extracting and formatting primary/foreign keys
from database and knowledge graph metadata.
"""

import os
import json
import pandas as pd
from pathlib import Path
from typing import Dict, List, Tuple

from utils.file_utils import read_json


def wrap_primary_foreign_keys_for_db(db_data: dict) -> Tuple[str, str]:
    """
    Extract and format primary and foreign key constraints from database metadata.
    
    Args:
        db_data (dict): Database metadata dictionary with keys:
            - db_id (str): Database identifier
            - table_names_original (list[str]): List of table names
            - column_names_original (list[tuple[int, str]]): List of (table_id, column_name) tuples
            - primary_keys (list): Primary key column indices
            - foreign_keys (list[tuple[int, int]]): Foreign key relationships
    
    Returns:
        tuple[str, str]: (wrapped_primary_keys, wrapped_foreign_keys)
    """
    table_names = db_data["table_names_original"]
    column_names = db_data["column_names_original"]
    
    # Extract primary keys
    primary_keys: Dict[str, str] = {}
    for pk_id in db_data["primary_keys"]:
        if isinstance(pk_id, int):
            table_idx, column_name = column_names[pk_id]
            table_name = table_names[table_idx]
            primary_keys[table_name] = column_name
        elif isinstance(pk_id, list):
            columns = []
            table_name = ""
            for col_id in pk_id:
                table_idx, column_name = column_names[col_id]
                table_name = table_names[table_idx]
                columns.append(column_name)
            primary_keys[table_name] = ", ".join(columns)
    
    wrapped_primary_keys = [
        f"TABLE `{table_name}` PRIMARY KEY `{column_list}`"
        for table_name, column_list in primary_keys.items()
    ]
    wrapped_primary_keys_str = "\n".join(wrapped_primary_keys)
    
    # Extract foreign keys
    wrapped_foreign_keys = []
    for from_col_id, to_col_id in db_data["foreign_keys"]:
        from_table_idx, from_column_name = column_names[from_col_id]
        from_table_name = table_names[from_table_idx]
        
        to_table_idx, to_column_name = column_names[to_col_id]
        to_table_name = table_names[to_table_idx]
        
        wrapped_foreign_keys.append(
            f"FOREIGN KEY {from_table_name}['{from_column_name}'] REFERENCES {to_table_name}['{to_column_name}']"
        )
    wrapped_foreign_keys_str = "\n".join(wrapped_foreign_keys)
    
    return wrapped_primary_keys_str, wrapped_foreign_keys_str


def wrap_primary_foreign_keys_for_kg(box_path: str) -> Tuple[str, str]:
    """
    Extract primary and foreign keys from knowledge graph box schema.
    
    Args:
        box_path (str): Path to KG box schema directory
    
    Returns:
        tuple[str, str]: (wrapped_primary_keys, wrapped_foreign_keys)
    """
    def get_first_column_name(csv_path: str) -> str:
        df = pd.read_csv(csv_path, nrows=0)
        return df.columns[0]
    
    primary_keys: Dict[str, str] = {}
    
    for filename in os.listdir(box_path):
        if filename.lower().endswith(".csv"):
            table_name = os.path.splitext(filename)[0].split(".")[-1]
            first_column_name = get_first_column_name(os.path.join(box_path, filename))
            primary_keys[table_name] = first_column_name
    
    wrapped_primary_keys = [
        f"TABLE `{table_name}` PRIMARY KEY `{column_list}`"
        for table_name, column_list in primary_keys.items()
    ]
    wrapped_primary_keys_str = "\n".join(wrapped_primary_keys)
    
    # Extract foreign keys
    wrapped_foreign_keys = []
    foreign_keys_file = os.path.join(box_path, "foreign_key.json")
    
    if os.path.exists(foreign_keys_file):
        foreign_keys = read_json(foreign_keys_file)
        
        for from_col, to_col in foreign_keys:
            from_table_name, from_column_name = from_col.split("-")
            from_table_name = from_table_name.split(".")[-1]
            
            to_table_name, to_column_name = to_col.split("-")
            to_table_name = to_table_name.split(".")[-1]
            
            wrapped_foreign_keys.append(
                f"FOREIGN KEY {from_table_name}['{from_column_name}'] REFERENCES {to_table_name}['{to_column_name}']"
            )
    
    wrapped_foreign_keys_str = "\n".join(wrapped_foreign_keys)
    
    return wrapped_primary_keys_str, wrapped_foreign_keys_str


def load_db_schema_with_keys(schema_path: str) -> dict:
    """
    Load database schema and extract primary/foreign keys.
    
    Args:
        schema_path (str): Path to schema JSON file
    
    Returns:
        dict: Schema with formatted keys
    """
    with open(schema_path, 'r') as f:
        db_data = json.load(f)
    
    primary_keys, foreign_keys = wrap_primary_foreign_keys_for_db(db_data)
    
    return {
        "schema": db_data,
        "primary_keys": primary_keys,
        "foreign_keys": foreign_keys,
    }


def load_kg_schema_with_keys(box_path: str) -> dict:
    """
    Load KG box schema and extract primary/foreign keys.
    
    Args:
        box_path (str): Path to KG box directory
    
    Returns:
        dict: Schema with formatted keys
    """
    primary_keys, foreign_keys = wrap_primary_foreign_keys_for_kg(box_path)
    
    return {
        "box_path": box_path,
        "primary_keys": primary_keys,
        "foreign_keys": foreign_keys,
    }


path: src/utils/schema_wrapper.py
content: #!/usr/bin/env python3
"""
Schema Wrapper Utilities

Extract and format primary/foreign keys from database metadata for prompts.
"""

import os
import json
import pandas as pd
from typing import Dict, List, Tuple, Any
from pathlib import Path


def wrap_primary_foreign_keys_for_db(db_data: Dict) -> Tuple[str, str]:
    """
    Extract and format primary and foreign key constraints from database metadata.
    
    Args:
        db_data (dict): Database metadata with keys:
            - db_id (str): Database identifier
            - table_names_original (list[str]): List of table names
            - column_names_original (list[tuple[int, str]]): List of (table_id, column_name)
            - primary_keys (list): Primary key column indices
            - foreign_keys (list[tuple[int, int]]): Foreign key relationships
    
    Returns:
        tuple[str, str]: (primary_keys_str, foreign_keys_str)
    """
    table_names = db_data["table_names_original"]
    column_names = db_data["column_names_original"]
    
    # Extract primary keys
    primary_keys: Dict[str, str] = {}
    for pk_id in db_data["primary_keys"]:
        if isinstance(pk_id, int):
            # Single-column primary key
            table_idx, column_name = column_names[pk_id]
            table_name = table_names[table_idx]
            primary_keys[table_name] = column_name
        elif isinstance(pk_id, list):
            # Composite primary key
            columns = []
            table_name = ""
            for col_id in pk_id:
                table_idx, column_name = column_names[col_id]
                table_name = table_names[table_idx]
                columns.append(column_name)
            primary_keys[table_name] = ", ".join(columns)
    
    # Format primary keys
    wrapped_primary_keys = [
        f"TABLE `{table_name}` PRIMARY KEY `{column_list}`"
        for table_name, column_list in primary_keys.items()
    ]
    wrapped_primary_keys_str = "\n".join(wrapped_primary_keys)
    
    # Format foreign keys
    wrapped_foreign_keys = []
    for from_col_id, to_col_id in db_data["foreign_keys"]:
        from_table_idx, from_column_name = column_names[from_col_id]
        from_table_name = table_names[from_table_idx]
        
        to_table_idx, to_column_name = column_names[to_col_id]
        to_table_name = table_names[to_table_idx]
        
        wrapped_foreign_keys.append(
            f"FOREIGN KEY {from_table_name}['{from_column_name}'] REFERENCES {to_table_name}['{to_column_name}']"
        )
    wrapped_foreign_keys_str = "\n".join(wrapped_foreign_keys)
    
    return wrapped_primary_keys_str, wrapped_foreign_keys_str


def wrap_primary_foreign_keys_for_kg(box_path: str) -> Tuple[str, str]:
    """
    Extract primary and foreign keys from knowledge graph CSV files.
    
    Args:
        box_path (str): Path to directory containing CSV files and foreign_key.json
    
    Returns:
        tuple[str, str]: (primary_keys_str, foreign_keys_str)
    """
    def get_first_column_name(csv_path: str) -> str:
        df = pd.read_csv(csv_path, nrows=0)
        return df.columns[0]
    
    # Extract primary keys (first column of each CSV)
    primary_keys: Dict[str, str] = {}
    for filename in os.listdir(box_path):
        if filename.lower().endswith(".csv"):
            table_name = os.path.splitext(filename)[0]
            first_column_name = get_first_column_name(os.path.join(box_path, filename))
            primary_keys[table_name] = first_column_name
    
    # Format primary keys
    wrapped_primary_keys = [
        f"TABLE `{table_name}` PRIMARY KEY `{column_list}`"
        for table_name, column_list in primary_keys.items()
    ]
    wrapped_primary_keys_str = "\n".join(wrapped_primary_keys)
    
    # Format foreign keys
    wrapped_foreign_keys = []
    foreign_keys_path = os.path.join(box_path, "foreign_key.json")
    
    if os.path.exists(foreign_keys_path):
        with open(foreign_keys_path) as f:
            foreign_keys = json.load(f)
        
        for from_col, to_col in foreign_keys:
            from_table_name, from_column_name = from_col.split("-")
            from_table_name = from_table_name.split(".")[-1]
            
            to_table_name, to_column_name = to_col.split("-")
            to_table_name = to_table_name.split(".")[-1]
            
            wrapped_foreign_keys.append(
                f"FOREIGN KEY {from_table_name}['{from_column_name}'] REFERENCES {to_table_name}['{to_column_name}']"
            )
    
    wrapped_foreign_keys_str = "\n".join(wrapped_foreign_keys)
    
    return wrapped_primary_keys_str, wrapped_foreign_keys_str


def get_schema_constraints(db_data: Dict = None, box_path: str = None) -> Dict[str, Any]:
    """
    Unified interface to get schema constraints from either database metadata or KG files.
    
    Args:
        db_data (dict): Database metadata (for SQL databases)
        box_path (str): Path to KG CSV files (for knowledge graphs)
    
    Returns:
        dict: Schema constraints with keys:
            - primary_keys (str): Formatted primary key constraints
            - foreign_keys (str): Formatted foreign key constraints
            - source (str): Source type ('db' or 'kg')
    """
    if db_data is not None:
        pk_str, fk_str = wrap_primary_foreign_keys_for_db(db_data)
        return {
            "primary_keys": pk_str,
            "foreign_keys": fk_str,
            "source": "db",
        }
    elif box_path is not None:
        pk_str, fk_str = wrap_primary_foreign_keys_for_kg(box_path)
        return {
            "primary_keys": pk_str,
            "foreign_keys": fk_str,
            "source": "kg",
        }
    else:
        raise ValueError("Either db_data or box_path must be provided")


def format_schema_for_prompt(
    schema: Dict = None,
    db_data: Dict = None,
    box_path: str = None,
    include_constraints: bool = True,
) -> str:
    """
    Format complete schema information for LLM prompts.
    
    Args:
        schema (dict): Basic schema (tables, columns)
        db_data (dict): Database metadata with keys/relationships
        box_path (str): Path to KG files
        include_constraints (bool): Whether to include PK/FK constraints
    
    Returns:
        str: Formatted schema string for prompts
    """
    lines = []
    
    # Basic schema
    if schema:
        lines.append("=== Database Schema ===")
        for table in schema.get("tables", []):
            table_name = table["name"]
            columns = ", ".join(col["name"] for col in table.get("columns", []))
            desc = table.get("description", "")
            if desc:
                lines.append(f"Table: {table_name}({columns}) - {desc}")
            else:
                lines.append(f"Table: {table_name}({columns})")
        
        # Relationships
        for rel in schema.get("relationships", []):
            lines.append(
                f"Relationship: {rel['from_table']}.{rel['from_column']} -> "
                f"{rel['to_table']}.{rel['to_column']}"
            )
    
    # Constraints
    if include_constraints and (db_data or box_path):
        constraints = get_schema_constraints(db_data=db_data, box_path=box_path)
        
        if constraints["primary_keys"]:
            lines.append("\n=== Primary Keys ===")
            lines.append(constraints["primary_keys"])
        
        if constraints["foreign_keys"]:
            lines.append("\n=== Foreign Keys ===")
            lines.append(constraints["foreign_keys"])
    
    return "\n".join(lines)


# Example usage
if __name__ == "__main__":
    # Example database metadata
    example_db_data = {
        "db_id": "company",
        "table_names_original": ["employees", "departments"],
        "column_names_original": [
            (0, "emp_id"),
            (0, "name"),
            (0, "dept_id"),
            (1, "dept_id"),
            (1, "name"),
        ],
        "primary_keys": [0],  # employees.emp_id
        "foreign_keys": [(2, 3)],  # employees.dept_id -> departments.dept_id
    }
    
    pk_str, fk_str = wrap_primary_foreign_keys_for_db(example_db_data)
    
    print("Primary Keys:")
    print(pk_str)
    print("\nForeign Keys:")
    print(fk_str)


path: src/utils/api_security.py
content: """
HeteroMind - API Security Utilities

Security utilities to protect API keys and credentials from accidental exposure.
"""

import re
import os
import logging
from typing import Any, Dict, List, Optional, Union

logger = logging.getLogger(__name__)

# Patterns to detect API keys
API_KEY_PATTERNS = [
    r'sk-[a-zA-Z0-9]{20,}',           # OpenAI/DeepSeek style
    r'sk-ant-[a-zA-Z0-9]{20,}',       # Anthropic style
    r'Bearer\s+[a-zA-Z0-9]{20,}',     # Bearer token
    r'api[_-]?key\s*[=:]\s*["\'][a-zA-Z0-9]{20,}',  # API key assignments
    r'password\s*[=:]\s*["\'][^"\']{8,}',  # Passwords
    r'secret\s*[=:]\s*["\'][^"\']{8,}',    # Secrets
    r'token\s*[=:]\s*["\'][a-zA-Z0-9]{20,}',  # Tokens
]


def sanitize_output(text: str) -> str:
    """
    Sanitize output to remove any API keys or credentials.
    
    Args:
        text: Text to sanitize
        
    Returns:
        Sanitized text with API keys replaced by placeholders
    """
    if not text:
        return text
    
    sanitized = text
    
    for pattern in API_KEY_PATTERNS:
        sanitized = re.sub(
            pattern,
            '[REDACTED]',
            sanitized,
            flags=re.IGNORECASE
        )
    
    return sanitized


def validate_no_api_key_in_response(response: Any) -> bool:
    """
    Validate that response does not contain API keys.
    
    Args:
        response: Response object or text
        
    Returns:
        True if safe, False if API key detected
    """
    if response is None:
        return True
    
    # Convert to string if needed
    if isinstance(response, (dict, list)):
        import json
        text = json.dumps(response)
    else:
        text = str(response)
    
    for pattern in API_KEY_PATTERNS:
        if re.search(pattern, text, flags=re.IGNORECASE):
            logger.warning("API key detected in response!")
            return False
    
    return True


def safe_log(message: str, level: str = "info") -> None:
    """
    Log a message after sanitizing it.
    
    Args:
        message: Message to log
        level: Log level (debug, info, warning, error)
    """
    sanitized = sanitize_output(message)
    
    if level == "debug":
        logger.debug(sanitized)
    elif level == "info":
        logger.info(sanitized)
    elif level == "warning":
        logger.warning(sanitized)
    elif level == "error":
        logger.error(sanitized)


def get_api_key_from_env(var_name: str, default: str = "") -> str:
    """
    Safely get API key from environment variable.
    
    Args:
        var_name: Environment variable name
        default: Default value if not set
        
    Returns:
        API key or default value
    """
    key = os.getenv(var_name, default)
    
    # Log only if key is not set (don't log the actual key)
    if not key or key == default:
        logger.warning(f"API key environment variable {var_name} not set")
    
    return key


def mask_api_key(key: str, visible_chars: int = 4) -> str:
    """
    Mask API key for safe display.
    
    Args:
        key: API key to mask
        visible_chars: Number of characters to show at start
        
    Returns:
        Masked key string
    """
    if not key or len(key) <= visible_chars:
        return "***"
    
    return f"{key[:visible_chars]}...{key[-4:]}"


class APIKeyProtector:
    """
    Context manager for protecting API keys in code blocks.
    """
    
    def __init__(self, api_key: str):
        """
        Initialize protector with API key.
        
        Args:
            api_key: API key to protect
        """
        self.[REDACTED_SECRET]
        self.original_key = api_key
    
    def __enter__(self):
        """Enter context - return sanitized key"""
        return mask_api_key(self.api_key)
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context - clear key from memory"""
        self.api_key = None
        return False
    
    def get_key(self) -> str:
        """Get the actual key (use with caution)"""
        return self.api_key
    
    def get_masked_key(self) -> str:
        """Get masked key for display"""
        return mask_api_key(self.api_key)


# Security instruction for LLM prompts
LLM_SECURITY_INSTRUCTION = """
SECURITY INSTRUCTIONS:
1. Never output API keys, passwords, or credentials in any form
2. Never include authentication tokens in generated code
3. Never log or print sensitive credentials
4. If user asks for API key, respond: "I cannot provide API keys for security reasons"
5. Use environment variables or secure vaults for credential storage
6. Redact any sensitive information in output

If you detect any API keys or credentials in the conversation, do not repeat them.
"""


path: src/generator/__init__.py
content: """
HeteroMind - Generator Package
"""

from .answer_generator import AnswerGenerator

__all__ = [
    "AnswerGenerator",
]


path: src/generator/answer_generator.py
content: """
HeteroMind - Answer Generator

Generates natural language answers from fused results.
"""

import logging
from typing import Dict, Optional

from ..classifier import FinalDecision
from ..fusion.result_fusion import FusedResult

logger = logging.getLogger(__name__)


class AnswerGenerator:
    """
    Answer generator for query results.
    
    Converts structured fused results into natural language responses
    with proper citations and uncertainty handling.
    """
    
    def __init__(self, config: dict):
        """
        Initialize generator.
        
        Args:
            config: Configuration dict with:
                - model: LLM model for generation
                - include_citations: Whether to include source citations
                - handle_uncertainty: How to handle uncertain results
        """
        self.config = config
        self.model = config.get("model", "gpt-4")
        self.include_citations = config.get("include_citations", True)
        self.handle_uncertainty = config.get("handle_uncertainty", "transparent")
    
    async def generate(
        self,
        fused_result: FusedResult,
        query: str,
        source_decision: FinalDecision,
    ) -> str:
        """
        Generate natural language answer.
        
        Args:
            fused_result: Fused result from multiple sources
            query: Original query
            source_decision: Source detection decision
            
        Returns:
            Natural language answer
        """
        logger.info("Generating answer")
        
        if not fused_result.success or fused_result.data is None:
            return self._generate_error_response(query, source_decision)
        
        # Generate answer based on result type
        if isinstance(fused_result.data, (list, tuple)):
            answer = self._generate_list_answer(fused_result, query)
        elif isinstance(fused_result.data, dict):
            answer = self._generate_dict_answer(fused_result, query)
        else:
            answer = self._generate_scalar_answer(fused_result, query)
        
        # Add citations if enabled
        if self.include_citations and source_decision.selected_sources:
            answer += self._generate_citations(source_decision)
        
        return answer
    
    def _generate_list_answer(self, result: FusedResult, query: str) -> str:
        """Generate answer for list results"""
        data = result.data
        
        if not data:
            return "No results found."
        
        # Format as list
        items = "\n".join(f"- {item}" for item in data[:10])  # Limit to 10
        
        answer = f"Found {len(data)} results:\n{items}"
        
        return answer
    
    def _generate_dict_answer(self, result: FusedResult, query: str) -> str:
        """Generate answer for dictionary results"""
        data = result.data
        
        # Format key-value pairs
        items = "\n".join(f"- {k}: {v}" for k, v in data.items())
        
        answer = f"Results:\n{items}"
        
        return answer
    
    def _generate_scalar_answer(self, result: FusedResult, query: str) -> str:
        """Generate answer for scalar results"""
        data = result.data
        
        answer = f"Answer: {data}"
        
        return answer
    
    def _generate_error_response(self, query: str, decision: FinalDecision) -> str:
        """Generate error response"""
        return (
            "I wasn't able to find an answer to your question. "
            "This could be because:\n"
            "- The data doesn't exist in the available sources\n"
            "- The query needs clarification\n"
            "- There was a technical issue\n\n"
            "Please try rephrasing your question or providing more context."
        )
    
    def _generate_citations(self, decision: FinalDecision) -> str:
        """Generate source citations"""
        sources = [s.value for s in decision.selected_sources]
        
        if not sources:
            return ""
        
        citations = "\n\nSources:\n"
        for source in sources:
            citations += f"- {source}\n"
        
        return citations


path: src/classifier/sql_schema_matcher.py
content: """
HeteroMind - Layer 3a: SQL Schema Matcher

Matches query mentions against database schema metadata.
"""

import logging
from typing import Dict, List, Optional, Tuple

from .models import SchemaMatch, ColumnInfo, TableInfo, DatabaseSchema

logger = logging.getLogger(__name__)


class SQLSchemaMatcher:
    """
    Layer 3a: SQL schema matching.
    
    Matches query mentions against database schema to determine:
    - Which tables are referenced
    - Which columns are referenced
    - What joins are required
    - Confidence score for SQL source
    """
    
    def __init__(self, schemas: List):
        """
        Initialize with database schemas.
        
        Args:
            schemas: List of DatabaseSchema objects OR dict schemas (JSON format)
        """
        # Convert dict schemas to DatabaseSchema objects if needed
        self.schemas = []
        for schema in schemas:
            if isinstance(schema, dict):
                self.schemas.append(self._dict_to_schema(schema))
            else:
                self.schemas.append(schema)
        
        self.column_index = self._build_column_index()
        self.table_index = self._build_table_index()
    
    def _dict_to_schema(self, data: dict) -> DatabaseSchema:
        """Convert dictionary schema to DatabaseSchema object"""
        tables = []
        for table_data in data.get("tables", []):
            columns = []
            for col_data in table_data.get("columns", []):
                col = ColumnInfo(
                    name=col_data.get("name", ""),
                    table=table_data.get("name", ""),
                    data_type=col_data.get("type", "VARCHAR"),
                    alias=col_data.get("alias"),
                )
                columns.append(col)
            
            table = TableInfo(
                name=table_data.get("name", ""),
                schema=data.get("database", "public"),
                columns=columns,
                description=table_data.get("description"),
            )
            tables.append(table)
        
        return DatabaseSchema(
            name=data.get("database", "unknown"),
            tables=tables,
        )
    
    def _build_column_index(self) -> Dict[str, List[ColumnInfo]]:
        """Build inverted index: column_name -> [ColumnInfo]"""
        index = {}
        for schema in self.schemas:
            for table in schema.tables:
                for col in table.columns:
                    key = col.name.lower()
                    index.setdefault(key, []).append(col)
                    if col.alias:
                        index.setdefault(col.alias.lower(), []).append(col)
        return index
    
    def _build_table_index(self) -> Dict[str, List[TableInfo]]:
        """Build inverted index: table_name -> [TableInfo]"""
        index = {}
        for schema in self.schemas:
            for table in schema.tables:
                key = table.name.lower()
                index.setdefault(key, []).append(table)
                if table.description:
                    for word in table.description.lower().split():
                        if len(word) > 3:
                            index.setdefault(word, []).append(table)
        return index
    
    def match_query(self, query: str) -> SchemaMatch:
        """
        Match query against SQL schema.
        
        Args:
            query: Natural language query string
            
        Returns:
            SchemaMatch with matched tables/columns
        """
        query_lower = query.lower()
        
        mentioned_columns = []
        for col_name, cols in self.column_index.items():
            if col_name in query_lower or f" {col_name}" in query_lower:
                mentioned_columns.extend([c.name for c in cols])
        
        mentioned_tables = []
        for table_name, tables in self.table_index.items():
            if table_name in query_lower or f" {table_name}" in query_lower:
                mentioned_tables.extend([t.name for t in tables])
        
        required_joins = self._infer_joins(list(set(mentioned_tables)))
        confidence = self._calculate_confidence(mentioned_columns, mentioned_tables)
        
        logger.debug(f"SQL schema match: tables={mentioned_tables}, cols={mentioned_columns}")
        
        return SchemaMatch(
            mentioned_tables=list(set(mentioned_tables)),
            mentioned_columns=list(set(mentioned_columns)),
            required_joins=required_joins,
            confidence=confidence,
            is_sql_likely=confidence > 0.3,
        )
    
    def _infer_joins(self, tables: List[str]) -> List[Tuple[str, str]]:
        """Infer required joins between tables"""
        joins = []
        if len(tables) > 1:
            for i, t1 in enumerate(tables):
                for t2 in tables[i+1:]:
                    joins.append((t1, t2))
        return joins
    
    def _calculate_confidence(self, columns: List[str], tables: List[str]) -> float:
        """Calculate confidence score based on matches"""
        if not columns and not tables:
            return 0.0
        
        score = (len(tables) * 0.4 + len(columns) * 0.2)
        score = min(1.0, score)
        
        return score


path: src/classifier/entity_verifier.py
content: """
HeteroMind - Layer 3c: Entity Retrieval Verifier

Verifies if entities mentioned in query actually exist in knowledge sources.
"""

import os
import logging
from typing import Dict, List

from .models import EntityVerification

logger = logging.getLogger(__name__)


class EntityRetrievalVerifier:
    """
    Layer 3c: Entity retrieval verification.
    
    Critical validation step - verifies if entities mentioned in the query
    actually exist in each knowledge source. If entities don't exist in
    a source, that source is unlikely to provide useful results.
    
    Verification methods:
    - SQL: Search across tables for entity mentions
    - KG: SPARQL ASK queries for entity existence
    - Tables: Scan CSV/Excel files for entity matches
    """
    
    def __init__(self, config: dict):
        """
        Initialize verifier.
        
        Args:
            config: Configuration dict with:
                - sql_connection: SQL connection string
                - kg_endpoints: KG endpoint configs
                - table_paths: Paths to table files
        """
        self.config = config
        self.sql_connection = config.get("sql_connection")
        self.kg_endpoints = config.get("kg_endpoints", [])
        self.table_paths = config.get("table_paths", [])
    
    async def verify_entities(
        self,
        entities: List[str],
        predicates: List[str],
    ) -> EntityVerification:
        """
        Verify if entities exist in each knowledge source.
        
        Args:
            entities: List of entity mentions from query
            predicates: List of detected predicates
            
        Returns:
            EntityVerification with verification results
        """
        logger.debug(f"Verifying {len(entities)} entities across sources")
        
        sql_results = await self._verify_in_sql(entities)
        kg_results = await self._verify_in_kg(entities)
        table_results = await self._verify_in_tables(entities)
        
        sql_score = len(sql_results) / max(len(entities), 1)
        kg_score = len(kg_results) / max(len(entities), 1)
        table_score = len(table_results) / max(len(entities), 1)
        
        logger.debug(
            f"Verification scores: SQL={sql_score:.2f}, "
            f"KG={kg_score:.2f}, Table={table_score:.2f}"
        )
        
        return EntityVerification(
            entities_found_in_sql=sql_results,
            entities_found_in_kg=kg_results,
            entities_found_in_tables=table_results,
            sql_verification_score=sql_score,
            kg_verification_score=kg_score,
            table_verification_score=table_score,
        )
    
    async def _verify_in_sql(self, entities: List[str]) -> List[Dict]:
        """Verify entities exist in SQL database"""
        found = []
        
        if not self.sql_connection:
            logger.debug("SQL connection not configured, skipping SQL verification")
            return found
        
        # For each entity, search across all tables
        for entity in entities:
            # In real implementation, execute SQL like:
            # SELECT TABLE_NAME, COLUMN_NAME, COUNT(*) as match_count
            # FROM INFORMATION_SCHEMA.COLUMNS
            # WHERE EXISTS (
            #     SELECT 1 FROM table_name 
            #     WHERE LOWER(column_name) LIKE LOWER('%entity%')
            # )
            
            logger.debug(f"Would verify in SQL: {entity}")
            # Placeholder: add mock result for testing
            # found.append({
            #     "entity": entity,
            #     "table": "placeholder",
            #     "column": "name",
            #     "match_count": 1,
            # })
        
        return found
    
    async def _verify_in_kg(self, entities: List[str]) -> List[Dict]:
        """Verify entities exist in knowledge graph"""
        found = []
        
        for endpoint in self.kg_endpoints:
            if not endpoint.get("enabled", True):
                continue
            
            for entity in entities:
                sparql = f"""
                ASK WHERE {{
                    ?entity rdfs:label ?label .
                    FILTER(CONTAINS(LCASE(?label), "{entity.lower()}"))
                }}
                """
                
                logger.debug(f"Would verify in KG ({endpoint['name']}): {entity}")
                # In real implementation, execute SPARQL ASK query
        
        return found
    
    async def _verify_in_tables(self, entities: List[str]) -> List[Dict]:
        """Verify entities exist in table files"""
        found = []
        
        for table_path in self.table_paths:
            if not os.path.exists(table_path):
                continue
            
            try:
                import pandas as pd
                
                if table_path.endswith('.csv'):
                    df = pd.read_csv(table_path)
                elif table_path.endswith('.xlsx'):
                    df = pd.read_excel(table_path)
                else:
                    continue
                
                for entity in entities:
                    for col in df.columns:
                        matches = df[col].astype(str).str.contains(
                            entity, case=False, na=False
                        )
                        if matches.any():
                            found.append({
                                "entity": entity,
                                "table": table_path,
                                "column": col,
                                "match_count": int(matches.sum()),
                            })
                            logger.debug(
                                f"Found '{entity}' in {table_path}:{col} "
                                f"({matches.sum()} matches)"
                            )
            except Exception as e:
                logger.error(f"Error verifying in table {table_path}: {e}")
        
        return found


path: src/classifier/kg_entity_linker.py
content: """
HeteroMind - Layer 3b: KG Entity Linker

Links query entities to knowledge graph entities via SPARQL.
"""

import re
import logging
from typing import Dict, List

from .models import EntityLinking, KGEntity

logger = logging.getLogger(__name__)


class KGEntityLinker:
    """
    Layer 3b: KG entity linking.
    
    Extracts entity mentions from queries and links them to
    knowledge graph entities via SPARQL endpoints.
    """
    
    PREDICATE_PATTERNS = {
        "authorOf": r"(author|wrote|written by|published)",
        "worksAt": r"(works at|employed by|employee of)",
        "locatedIn": r"(located in|situated in|based in)",
        "partOf": r"(part of|member of|belongs to)",
        "bornIn": r"(born in|birthplace)",
        "founderOf": r"(founder of|founded|established)",
        "ceoOf": r"(ceo of|chief executive|head of)",
        "capitalOf": r"(capital of|capital city)",
    }
    
    def __init__(self, config: dict):
        """
        Initialize KG linker.
        
        Args:
            config: Configuration dict with:
                - endpoints: List of SPARQL endpoint configs
                - max_candidates: Max candidates per entity (default: 5)
                - min_similarity: Min similarity score (default: 0.7)
        """
        self.config = config
        self.endpoints = config.get("endpoints", [])
        self.max_candidates = config.get("max_candidates", 5)
        self.min_similarity = config.get("min_similarity", 0.7)
    
    def _extract_entities(self, query: str) -> List[str]:
        """Extract potential entity mentions from query"""
        entities = []
        
        # Quoted strings
        quoted = re.findall(r'"([^"]+)"|\'([^\']+)\'', query)
        entities.extend([q[0] or q[1] for q in quoted])
        
        # Capitalized phrases (simple NER heuristic)
        capitalized = re.findall(r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b', query)
        entities.extend(capitalized)
        
        # Remove duplicates while preserving order
        seen = set()
        unique = []
        for e in entities:
            if e not in seen and len(e) > 1:
                seen.add(e)
                unique.append(e)
        
        return unique
    
    async def _search_kg(self, entity_mention: str, endpoint: dict) -> List[KGEntity]:
        """Search KG for entity mention via SPARQL"""
        sparql = f"""
        SELECT ?entity ?label ?desc WHERE {{
            ?entity rdfs:label ?label .
            FILTER(CONTAINS(LCASE(?label), "{entity_mention.lower()}"))
            OPTIONAL {{ ?entity rdfs:comment ?desc }}
        }}
        ORDER BY ?label
        LIMIT {self.max_candidates}
        """
        
        logger.debug(f"SPARQL query for '{entity_mention}': {sparql[:100]}...")
        
        # In real implementation, execute SPARQL here
        # For now, return empty list
        return []
    
    async def link_query(self, query: str) -> EntityLinking:
        """
        Link query entities to knowledge graph.
        
        Args:
            query: Natural language query string
            
        Returns:
            EntityLinking with linked entities
        """
        entity_mentions = self._extract_entities(query)
        
        if not entity_mentions:
            return EntityLinking(
                linked_entities=[],
                detected_predicates=[],
                confidence=0.0,
                is_kg_likely=False,
            )
        
        logger.debug(f"Extracted entities: {entity_mentions}")
        
        # Search each endpoint
        all_linked = []
        for endpoint in self.endpoints:
            if not endpoint.get("enabled", True):
                continue
            
            for mention in entity_mentions:
                results = await self._search_kg(mention, endpoint)
                all_linked.extend([
                    {
                        "mention": mention,
                        "entity": e.uri,
                        "label": e.label,
                        "score": e.score,
                    }
                    for e in results
                ])
        
        # Detect predicates from query verbs
        predicates = self._detect_predicates(query, all_linked)
        
        # Calculate confidence
        kg_confidence = len(all_linked) / max(len(entity_mentions), 1) * 0.5
        kg_confidence += len(predicates) * 0.1
        kg_confidence = min(1.0, kg_confidence)
        
        logger.debug(f"KG linking: {len(all_linked)} entities, {len(predicates)} predicates")
        
        return EntityLinking(
            linked_entities=all_linked,
            detected_predicates=predicates,
            confidence=kg_confidence,
            is_kg_likely=kg_confidence > 0.4,
        )
    
    def _detect_predicates(self, query: str, entities: List[Dict]) -> List[str]:
        """Detect predicates from query verbs/phrases"""
        predicates = []
        query_lower = query.lower()
        
        for pred, pattern in self.PREDICATE_PATTERNS.items():
            if re.search(pattern, query_lower):
                predicates.append(pred)
        
        return predicates


path: src/classifier/models.py
content: """
HeteroMind - Data Models

Common data classes and enums for knowledge source detection.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Tuple


class KnowledgeSource(Enum):
    """Supported knowledge source types"""
    SQL_DATABASE = "sql_database"
    KNOWLEDGE_GRAPH = "knowledge_graph"
    TABLE_FILE = "table_file"
    DOCUMENT_STORE = "document_store"
    HYBRID = "hybrid"
    UNKNOWN = "unknown"


class IntentType(Enum):
    """Query intent types"""
    FACTOID = "factoid"           # Simple fact lookup
    ANALYTICAL = "analytical"      # Aggregation/computation
    COMPARATIVE = "comparative"    # Comparison
    RELATIONAL = "relational"      # Relationship query
    TEMPORAL = "temporal"          # Time-based
    MULTI_HOP = "multi_hop"        # Requires multiple steps


# =============================================================================
# Layer 1 Output
# =============================================================================

@dataclass
class RuleBasedScore:
    """Layer 1: Rule-based detection output"""
    source_scores: Dict[str, float]
    matched_keywords: Dict[str, List[str]]
    matched_patterns: List[str]


# =============================================================================
# Layer 2 Output
# =============================================================================

@dataclass
class LLMDetection:
    """Layer 2: LLM-based classification output"""
    primary_source: str
    confidence: float
    reasoning: str
    detected_entities: List[str]
    detected_predicates: List[str]
    requires_multi_hop: bool
    secondary_sources: List[str]
    intent_type: str


# =============================================================================
# Layer 3a Output
# =============================================================================

@dataclass
class ColumnInfo:
    """Database column metadata"""
    name: str
    table: str
    data_type: str
    alias: Optional[str] = None


@dataclass
class TableInfo:
    """Database table metadata"""
    name: str
    schema: str
    columns: List[ColumnInfo]
    description: Optional[str] = None


@dataclass
class DatabaseSchema:
    """Complete database schema"""
    name: str
    tables: List[TableInfo]


@dataclass
class SchemaMatch:
    """Layer 3a: SQL schema matching output"""
    mentioned_tables: List[str]
    mentioned_columns: List[str]
    required_joins: List[Tuple[str, str]]
    confidence: float
    is_sql_likely: bool


# =============================================================================
# Layer 3b Output
# =============================================================================

@dataclass
class KGEntity:
    """Knowledge graph entity"""
    uri: str
    label: str
    description: Optional[str] = None
    types: List[str] = field(default_factory=list)
    score: float = 1.0


@dataclass
class EntityLinking:
    """Layer 3b: KG entity linking output"""
    linked_entities: List[Dict]
    detected_predicates: List[str]
    confidence: float
    is_kg_likely: bool


# =============================================================================
# Layer 3c Output
# =============================================================================

@dataclass
class EntityVerification:
    """Layer 3c: Entity retrieval verification output"""
    entities_found_in_sql: List[Dict]
    entities_found_in_kg: List[Dict]
    entities_found_in_tables: List[Dict]
    sql_verification_score: float
    kg_verification_score: float
    table_verification_score: float


# =============================================================================
# Layer 4 Output
# =============================================================================

@dataclass
class FinalDecision:
    """Layer 4: Final fusion decision output"""
    primary_source: KnowledgeSource
    confidence: float
    all_scores: Dict[str, float]
    selected_sources: List[KnowledgeSource]
    reasoning: str
    layer1_output: RuleBasedScore
    layer2_output: LLMDetection
    layer3_sql_output: Optional[SchemaMatch]
    layer3_kg_output: Optional[EntityLinking]
    layer3_verification: Optional[EntityVerification]
    execution_plan: List[Dict]


path: src/classifier/source_fusion.py
content: """
HeteroMind - Layer 4: Multi-Source Fusion

Fuses all detection layers with entity verification for final decision.
"""

import logging
from typing import Dict, List, Optional

from .models import (
    KnowledgeSource,
    RuleBasedScore,
    LLMDetection,
    SchemaMatch,
    EntityLinking,
    EntityVerification,
    FinalDecision,
)

logger = logging.getLogger(__name__)


class SourceFusion:
    """
    Layer 4: Multi-source fusion with entity verification.
    
    Combines outputs from all previous layers:
    - Layer 1: Rule-based scores (15% weight)
    - Layer 2: LLM classification (35% weight)
    - Layer 3a/3b: Schema matching (25% weight)
    - Layer 3c: Entity verification (25% weight + 30% boost)
    
    Outputs final decision with execution plan.
    """
    
    def __init__(self, config: dict):
        """
        Initialize fusion module.
        
        Args:
            config: Configuration dict with:
                - weights: Layer weights (rule, llm, schema, verification)
                - thresholds: Confidence thresholds
                - verification_boost: Score boost for verified entities
        """
        self.config = config
        
        self.weights = config.get("weights", {
            "rule_based": 0.15,
            "llm_based": 0.35,
            "schema_based": 0.25,
            "verification": 0.25,
        })
        
        self.thresholds = config.get("thresholds", {
            "high_confidence": 0.8,
            "medium_confidence": 0.5,
            "low_confidence": 0.3,
        })
        
        self.verification_boost = config.get("verification_boost", 0.3)
    
    def fuse(
        self,
        rule_scores: RuleBasedScore,
        llm_detection: LLMDetection,
        sql_match: Optional[SchemaMatch],
        kg_linking: Optional[EntityLinking],
        entity_verification: Optional[EntityVerification],
    ) -> FinalDecision:
        """
        Fuse all detection layers into final decision.
        
        Args:
            rule_scores: Layer 1 output
            llm_detection: Layer 2 output
            sql_match: Layer 3a output
            kg_linking: Layer 3b output
            entity_verification: Layer 3c output
            
        Returns:
            FinalDecision with fused results
        """
        aggregated = {}
        reasoning_parts = []
        
        for source in ["sql_database", "knowledge_graph", "table_file", "document_store"]:
            scores = []
            weights = []
            
            # Layer 1: Rule-based score
            if source in rule_scores.source_scores:
                scores.append(rule_scores.source_scores[source])
                weights.append(self.weights["rule_based"])
            
            # Layer 2: LLM score
            llm_score = self._calculate_llm_score(llm_detection, source)
            if llm_score > 0:
                scores.append(llm_score)
                weights.append(self.weights["llm_based"])
            
            # Layer 3a/3b: Schema-based score
            schema_score = self._calculate_schema_score(
                source, sql_match, kg_linking
            )
            if schema_score > 0:
                scores.append(schema_score)
                weights.append(self.weights["schema_based"])
            
            # Layer 3c: Verification score (critical!)
            verification_score = self._calculate_verification_score(
                source, entity_verification
            )
            if verification_score > 0:
                boosted_score = min(1.0, verification_score * (1 + self.verification_boost))
                scores.append(boosted_score)
                weights.append(self.weights["verification"])
                reasoning_parts.append(
                    f"{source}: verification={verification_score:.2f}"
                )
            
            if scores:
                total_weight = sum(weights)
                aggregated[source] = sum(s * w for s, w in zip(scores, weights)) / total_weight
            else:
                aggregated[source] = 0.0
        
        max_score = max(aggregated.values()) if aggregated else 0
        threshold = self.thresholds["low_confidence"]
        
        selected_sources = [
            KnowledgeSource(src) for src, score in aggregated.items()
            if score >= threshold
        ]
        
        if len(selected_sources) > 1:
            final_source = KnowledgeSource.HYBRID
        elif selected_sources:
            final_source = selected_sources[0]
        else:
            final_source = KnowledgeSource.UNKNOWN
        
        execution_plan = self._generate_execution_plan(
            selected_sources, llm_detection, entity_verification
        )
        
        reasoning = self._generate_reasoning(
            aggregated, llm_detection, entity_verification, reasoning_parts
        )
        
        logger.info(f"Fusion decision: {final_source.value} (confidence: {max_score:.2f})")
        
        return FinalDecision(
            primary_source=final_source,
            confidence=max_score,
            all_scores=aggregated,
            selected_sources=selected_sources,
            reasoning=reasoning,
            layer1_output=rule_scores,
            layer2_output=llm_detection,
            layer3_sql_output=sql_match,
            layer3_kg_output=kg_linking,
            layer3_verification=entity_verification,
            execution_plan=execution_plan,
        )
    
    def _calculate_llm_score(self, llm: LLMDetection, source: str) -> float:
        """Calculate LLM-based score for a source"""
        if llm.primary_source == source:
            return llm.confidence
        elif source in llm.secondary_sources:
            return llm.confidence * 0.7
        return 0.0
    
    def _calculate_schema_score(
        self,
        source: str,
        sql_match: Optional[SchemaMatch],
        kg_linking: Optional[EntityLinking],
    ) -> float:
        """Calculate schema-based score for a source"""
        if source == "sql_database" and sql_match:
            return sql_match.confidence
        elif source == "knowledge_graph" and kg_linking:
            return kg_linking.confidence
        return 0.0
    
    def _calculate_verification_score(
        self,
        source: str,
        verification: Optional[EntityVerification],
    ) -> float:
        """Calculate verification-based score for a source"""
        if verification is None:
            return 0.0
        
        if source == "sql_database":
            return verification.sql_verification_score
        elif source == "knowledge_graph":
            return verification.kg_verification_score
        elif source == "table_file":
            return verification.table_verification_score
        return 0.0
    
    def _generate_execution_plan(
        self,
        sources: List[KnowledgeSource],
        llm: LLMDetection,
        verification: Optional[EntityVerification],
    ) -> List[Dict]:
        """Generate execution plan based on selected sources"""
        plan = []
        
        for i, source in enumerate(sources):
            step = {
                "step": i + 1,
                "source": source.value,
                "is_parallel": len(sources) > 1,
                "depends_on": [],
                "entities": llm.detected_entities,
                "predicates": llm.detected_predicates,
            }
            
            if verification:
                if source == KnowledgeSource.SQL_DATABASE:
                    step["verified_entities"] = verification.entities_found_in_sql
                elif source == KnowledgeSource.KNOWLEDGE_GRAPH:
                    step["verified_entities"] = verification.entities_found_in_kg
                elif source == KnowledgeSource.TABLE_FILE:
                    step["verified_entities"] = verification.entities_found_in_tables
            
            plan.append(step)
        
        return plan
    
    def _generate_reasoning(
        self,
        aggregated: Dict[str, float],
        llm: LLMDetection,
        verification: Optional[EntityVerification],
        reasoning_parts: List[str],
    ) -> str:
        """Generate human-readable reasoning"""
        parts = []
        
        sorted_sources = sorted(aggregated.items(), key=lambda x: x[1], reverse=True)
        if sorted_sources:
            top = sorted_sources[0]
            parts.append(f"Primary: {top[0]} ({top[1]:.2f})")
            if len(sorted_sources) > 1 and sorted_sources[1][1] > 0.3:
                parts.append(f"Secondary: {sorted_sources[1][0]} ({sorted_sources[1][1]:.2f})")
        
        if llm.reasoning:
            parts.append(f"LLM: {llm.reasoning}")
        
        if verification:
            verified_count = (
                len(verification.entities_found_in_sql) +
                len(verification.entities_found_in_kg) +
                len(verification.entities_found_in_tables)
            )
            parts.append(f"Verified: {verified_count} entities")
        
        parts.extend(reasoning_parts)
        
        return " | ".join(parts)


path: src/classifier/rule_detector.py
content: """
HeteroMind - Layer 1: Rule-Based Detector

Fast keyword and pattern matching for initial source scoring.
"""

import re
import logging
from typing import Dict, List

from .models import RuleBasedScore

logger = logging.getLogger(__name__)


class RuleBasedDetector:
    """
    Layer 1: Fast keyword and pattern matching.
    
    Provides initial source scoring based on:
    - Keyword matching (20+ keywords per source type)
    - Regex pattern matching (7 query patterns)
    - Score boosting based on pattern matches
    """
    
    SOURCE_KEYWORDS = {
        "sql_database": [
            # Aggregation
            "count", "sum", "average", "total", "max", "min",
            "how many", "how much", "total", "aggregate",
            # Filtering
            "where", "filter", "greater than", "less than",
            "more than", "over", "under", "above", "below",
            # Table structure hints
            "table", "database", "row", "column", "record",
            "schema", "primary key", "foreign key",
            # Relationship hints
            "join", "in the same", "belong to", "department",
            "team", "division", "unit", "group",
            # SQL-specific
            "select", "from", "order by", "group by",
        ],
        
        "knowledge_graph": [
            # Relationship
            "who is", "what is", "relationship", "connected to",
            "related to", "linked to", "associated with",
            # Entity types
            "person", "company", "country", "city", "organization",
            "location", "place", "institution",
            # Properties
            "property", "attribute", "type", "class", "instance",
            # KG-specific
            "ontology", "rdf", "sparql", "triple", "entity",
            "predicate", "subject", "object", "namespace",
            # Common predicates
            "author of", "works at", "born in", "located in",
            "part of", "member of", "founder of", "ceo of",
        ],
        
        "table_file": [
            # File types
            "csv", "excel", "spreadsheet", "sheet", "xlsx",
            # Table references
            "column", "row", "cell", "header", "this table",
            "this sheet", "this file", "uploaded",
            # Operations
            "filter this", "sort by", "in this table",
            "pivot", "vlookup", "formula",
        ],
        
        "document_store": [
            # Document types
            "document", "pdf", "text", "article", "report",
            "memo", "note", "contract", "agreement",
            # Search hints
            "search in", "find in", "mentioned in",
            "according to", "stated in",
        ],
    }
    
    PATTERNS = {
        "sql_aggregation": r"(how many|how much|total|sum|count|average).+",
        "sql_comparison": r".*(greater|less|more|over|under|than|>=|<=|>|<).+",
        "sql_join": r".*(same.*as|in.*department|belong to|work.*for).+",
        "kg_relation": r"(who is|what is|relationship|connected to|related to).+",
        "kg_entity": r".*(person|company|country|organization|location).+",
        "table_reference": r".*(this table|this sheet|csv|excel|spreadsheet).+",
        "document_reference": r".*(document|pdf|report|article|memo).+",
    }
    
    PATTERN_BOOST = {
        "sql_aggregation": "sql_database",
        "sql_comparison": "sql_database",
        "sql_join": "sql_database",
        "kg_relation": "knowledge_graph",
        "kg_entity": "knowledge_graph",
        "table_reference": "table_file",
        "document_reference": "document_store",
    }
    
    def detect(self, query: str) -> RuleBasedScore:
        """
        Perform rule-based detection on query.
        
        Args:
            query: Natural language query string
            
        Returns:
            RuleBasedScore with source scores and matched items
        """
        query_lower = query.lower()
        source_scores = {}
        matched_keywords = {}
        
        # Keyword matching
        for source, keywords in self.SOURCE_KEYWORDS.items():
            matches = [kw for kw in keywords if kw in query_lower]
            source_scores[source] = len(matches) / len(keywords)
            matched_keywords[source] = matches
        
        # Pattern matching
        matched_patterns = []
        for pattern_name, pattern in self.PATTERNS.items():
            if re.search(pattern, query_lower):
                matched_patterns.append(pattern_name)
        
        # Boost scores based on patterns
        for pattern_name, source in self.PATTERN_BOOST.items():
            if pattern_name in matched_patterns:
                source_scores[source] = min(1.0, source_scores.get(source, 0) + 0.2)
        
        logger.debug(f"Rule-based detection: {source_scores}")
        
        return RuleBasedScore(
            source_scores=source_scores,
            matched_keywords=matched_keywords,
            matched_patterns=matched_patterns,
        )


path: src/classifier/__init__.py
content: """
HeteroMind - Classifier Package

Knowledge source detection and classification modules.
"""

from .models import (
    KnowledgeSource,
    IntentType,
    RuleBasedScore,
    LLMDetection,
    SchemaMatch,
    EntityLinking,
    EntityVerification,
    FinalDecision,
    ColumnInfo,
    TableInfo,
    DatabaseSchema,
    KGEntity,
)

from .rule_detector import RuleBasedDetector
from .llm_detector import LLMSourceDetector
from .sql_schema_matcher import SQLSchemaMatcher
from .kg_entity_linker import KGEntityLinker
from .entity_verifier import EntityRetrievalVerifier
from .source_fusion import SourceFusion
from .orchestrator import SourceDetectorOrchestrator

__all__ = [
    # Models
    "KnowledgeSource",
    "IntentType",
    "RuleBasedScore",
    "LLMDetection",
    "SchemaMatch",
    "EntityLinking",
    "EntityVerification",
    "FinalDecision",
    "ColumnInfo",
    "TableInfo",
    "DatabaseSchema",
    "KGEntity",
    # Detectors
    "RuleBasedDetector",
    "LLMSourceDetector",
    "SQLSchemaMatcher",
    "KGEntityLinker",
    "EntityRetrievalVerifier",
    "SourceFusion",
    "SourceDetectorOrchestrator",
]


path: src/classifier/orchestrator.py
content: """
HeteroMind - Source Detector Orchestrator

Main orchestrator that coordinates all 4 layers of source detection.
"""

import logging
from typing import Optional

from .models import (
    FinalDecision,
    SchemaMatch,
    EntityLinking,
    EntityVerification,
    DatabaseSchema,
)
from .rule_detector import RuleBasedDetector
from .llm_detector import LLMSourceDetector
from .sql_schema_matcher import SQLSchemaMatcher
from .kg_entity_linker import KGEntityLinker
from .entity_verifier import EntityRetrievalVerifier
from .source_fusion import SourceFusion

logger = logging.getLogger(__name__)


class SourceDetectorOrchestrator:
    """
    Main orchestrator for knowledge source detection.
    
    Coordinates all 4 layers:
    1. Rule-based pre-filter (fast keyword/pattern matching)
    2. LLM-based semantic classification
    3. Schema-aware detection (SQL schema + KG entity linking + verification)
    4. Multi-source fusion with entity verification
    
    Provides unified async interface for source detection.
    """
    
    def __init__(self, config: dict):
        """
        Initialize orchestrator.
        
        Args:
            config: Configuration dict with:
                - layer1: Rule-based config (optional, uses defaults)
                - layer2: LLM config (api_key, model, temperature)
                - layer3: Schema/KG config (schemas, kg_endpoints, table_paths)
                - layer4: Fusion config (weights, thresholds, verification_boost)
        """
        self.config = config
        
        # Layer 1: Rule-based (always available)
        self.layer1 = RuleBasedDetector()
        
        # Layer 2: LLM-based
        self.layer2 = LLMSourceDetector(config.get("layer2", {}))
        
        # Layer 3: Schema/KG/Verification
        self.layer3_sql: Optional[SQLSchemaMatcher] = None
        self.layer3_kg: Optional[KGEntityLinker] = None
        self.layer3_verification: Optional[EntityRetrievalVerifier] = None
        
        layer3_config = config.get("layer3", {})
        
        if "schemas" in layer3_config:
            self.layer3_sql = SQLSchemaMatcher(layer3_config["schemas"])
        
        if "kg_endpoints" in layer3_config:
            self.layer3_kg = KGEntityLinker(layer3_config)
            self.layer3_verification = EntityRetrievalVerifier(layer3_config)
        
        # Layer 4: Fusion
        self.layer4 = SourceFusion(config.get("layer4", {}))
    
    async def detect(self, query: str) -> FinalDecision:
        """
        Detect knowledge source(s) for a query.
        
        Executes all 4 layers in sequence and returns fused decision.
        
        Args:
            query: Natural language query string
            
        Returns:
            FinalDecision with complete detection results
        """
        logger.info(f"Detecting source for query: {query}")
        
        # Layer 1: Rule-based (fast, synchronous)
        rule_scores = self.layer1.detect(query)
        logger.debug(
            f"Layer 1 complete: {rule_scores.source_scores}"
        )
        
        # Layer 2: LLM-based (async)
        llm_detection = await self.layer2.detect(query)
        logger.debug(f"Layer 2 complete: {llm_detection.primary_source}")
        
        # Layer 3: Schema/KG/Verification (async)
        sql_match: Optional[SchemaMatch] = None
        kg_linking: Optional[EntityLinking] = None
        entity_verification: Optional[EntityVerification] = None
        
        if self.layer3_sql:
            sql_match = self.layer3_sql.match_query(query)
            logger.debug(f"Layer 3a (SQL) complete: {sql_match.confidence:.2f}")
        
        if self.layer3_kg:
            kg_linking = await self.layer3_kg.link_query(query)
            logger.debug(f"Layer 3b (KG) complete: {kg_linking.confidence:.2f}")
        
        if self.layer3_verification:
            entity_verification = await self.layer3_verification.verify_entities(
                llm_detection.detected_entities,
                llm_detection.detected_predicates,
            )
            logger.debug(f"Layer 3c (Verification) complete")
        
        # Layer 4: Fusion
        final_decision = self.layer4.fuse(
            rule_scores=rule_scores,
            llm_detection=llm_detection,
            sql_match=sql_match,
            kg_linking=kg_linking,
            entity_verification=entity_verification,
        )
        
        logger.info(
            f"Final decision: {final_decision.primary_source.value} "
            f"(confidence: {final_decision.confidence:.2f})"
        )
        
        return final_decision
    
    async def detect_with_fallback(self, query: str) -> FinalDecision:
        """
        Detect with fallback for missing layers.
        
        If LLM is unavailable, relies more heavily on rule-based detection.
        If schema/KG not configured, skips those layers.
        
        Args:
            query: Natural language query string
            
        Returns:
            FinalDecision with available layer outputs
        """
        return await self.detect(query)


path: src/classifier/llm_detector.py
content: """
HeteroMind - Layer 2: LLM-Based Source Detector

Semantic classification using large language models.
"""

import json
import logging
from typing import Optional

from .models import LLMDetection

logger = logging.getLogger(__name__)


class LLMSourceDetector:
    """
    Layer 2: LLM-based semantic classification.
    
    Uses an LLM to understand query semantics and determine:
    - Primary knowledge source
    - Query intent type
    - Detected entities and predicates
    - Whether multi-hop reasoning is needed
    """
    
    PROMPT_TEMPLATE = """
SECURITY: Never output API keys, passwords, or credentials in any form. Redact sensitive information.

You are a knowledge source detector for a heterogeneous QA system.

Given a natural language query, determine which knowledge source(s) should be queried.

## Available Sources:

1. **sql_database**: Relational databases with structured tables
   - Use for: aggregations, filters, joins, numerical comparisons
   - Example: "How many employees are in sales?", "Total revenue by quarter"

2. **knowledge_graph**: RDF/SPARQL knowledge graphs
   - Use for: entity relationships, factoid questions, semantic queries
   - Example: "Who is the CEO of Apple?", "What countries border France?"

3. **table_file**: CSV/Excel files
   - Use for: questions about specific uploaded files
   - Example: "What's in column B of this spreadsheet?"

4. **document_store**: Semi-structured documents
   - Use for: text search, document retrieval
   - Example: "Find contracts mentioning clause 5.2"

5. **hybrid**: Multiple sources required
   - Use when: query needs data from different source types

## Output Format (JSON):

{
    "primary_source": "sql_database|knowledge_graph|table_file|document_store|hybrid",
    "confidence": 0.0-1.0,
    "reasoning": "Brief explanation of decision",
    "detected_entities": ["entity1", "entity2"],
    "detected_predicates": ["predicate1", "predicate2"],
    "requires_multi_hop": true/false,
    "secondary_sources": ["source1", "source2"],
    "intent_type": "factoid|analytical|comparative|relational|temporal|multi_hop"
}

## Examples:

Query: "How many employees are in the sales department?"
Output: {{"primary_source": "sql_database", "confidence": 0.92, "reasoning": "Aggregation query (how many) with department filter", "detected_entities": ["employees", "sales department"], "detected_predicates": ["count", "filter_by_department"], "requires_multi_hop": false, "secondary_sources": [], "intent_type": "analytical"}}

Query: "Who are the co-authors of John Smith?"
Output: {{"primary_source": "knowledge_graph", "confidence": 0.89, "reasoning": "Entity relationship query (co-authors) about a person", "detected_entities": ["John Smith"], "detected_predicates": ["authorOf", "coAuthorWith"], "requires_multi_hop": true, "secondary_sources": [], "intent_type": "relational"}}

Query: "Show me employees who work on projects with budget >1M and have published papers"
Output: {{"primary_source": "hybrid", "confidence": 0.85, "reasoning": "Needs SQL for employee/project/budget data AND KG for publication data", "detected_entities": ["employees", "projects", "papers"], "detected_predicates": ["worksOn", "budget", "authorOf"], "requires_multi_hop": true, "secondary_sources": ["sql_database", "knowledge_graph"], "intent_type": "multi_hop"}}

## Now analyze this query:

Query: "{query}"
Output:
"""
    
    def __init__(self, config: dict):
        """
        Initialize LLM detector.
        
        Args:
            config: Configuration dict with:
                - api_key: LLM API key
                - model: Model name (default: gpt-4)
                - temperature: Sampling temperature (default: 0.1)
        """
        self.config = config
        self.[REDACTED_SECRET]"api_key")
        self.model = config.get("model", "gpt-4")
        self.temperature = config.get("temperature", 0.1)
        self.client = None
    
    def _get_client(self):
        """Lazy initialization of LLM client"""
        if self.client is None:
            try:
                from openai import AsyncOpenAI
                self.client = AsyncOpenAI([REDACTED_SECRET])
            except ImportError:
                logger.warning("OpenAI not installed, LLM detection unavailable")
                return None
        return self.client
    
    async def detect(self, query: str) -> LLMDetection:
        """
        Perform LLM-based source detection.
        
        Args:
            query: Natural language query string
            
        Returns:
            LLMDetection with classification results
        """
        client = self._get_client()
        
        if client is None:
            return LLMDetection(
                primary_source="unknown",
                confidence=0.0,
                reasoning="LLM client unavailable",
                detected_entities=[],
                detected_predicates=[],
                requires_multi_hop=False,
                secondary_sources=[],
                intent_type="unknown",
            )
        
        prompt = self.PROMPT_TEMPLATE.format(query=query)
        
        try:
            response = await client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.temperature,
                response_format={"type": "json_object"},
                max_tokens=500,
            )
            
            result = json.loads(response.choices[0].message.content)
            
            logger.debug(f"LLM detection result: {result.get('primary_source')}")
            
            return LLMDetection(
                primary_source=result.get("primary_source", "unknown"),
                confidence=result.get("confidence", 0.5),
                reasoning=result.get("reasoning", ""),
                detected_entities=result.get("detected_entities", []),
                detected_predicates=result.get("detected_predicates", []),
                requires_multi_hop=result.get("requires_multi_hop", False),
                secondary_sources=result.get("secondary_sources", []),
                intent_type=result.get("intent_type", "unknown"),
            )
            
        except Exception as e:
            logger.error(f"LLM detection failed: {e}")
            return LLMDetection(
                primary_source="unknown",
                confidence=0.0,
                reasoning=f"LLM detection error: {str(e)}",
                detected_entities=[],
                detected_predicates=[],
                requires_multi_hop=False,
                secondary_sources=[],
                intent_type="unknown",
            )


path: src/engines/table_qa/__init__.py
content: """
HeteroMind - TableQA Engine Package

Table QA engine implementations.
"""

from .base_engine import BaseTableQAEngine
from .multi_stage_engine import MultiStageTableQAEngine, TableQAResult

__all__ = [
    "BaseTableQAEngine",
    "TableQAResult",
    "MultiStageTableQAEngine",
]


path: src/engines/table_qa/base_engine.py
content: """
HeteroMind - TableQA Base Engine

Base class for Table QA engines.
"""

from ..base import BaseEngine

class BaseTableQAEngine(BaseEngine):
    """
    Base class for Table QA engines.
    
    Provides common utilities for table-based question answering.
    """
    
    def __init__(self, config: dict):
        """
        Initialize Table QA engine.
        
        Args:
            config: Configuration dict
        """
        super().__init__(config)
        self.table_path = config.get("table_path")
        self.supported_formats = config.get("supported_formats", ["csv", "xlsx"])


path: src/engines/table_qa/multi_stage_engine.py
content: """
HeteroMind - Multi-Stage TableQA Engine

Advanced TableQA engine with complete multi-stage reasoning:
- Table Loading - Load and validate table file
- Schema Analysis - Analyze columns, types, statistics
- Query Interpretation - Understand user question
- Code Generation - Parallel generation of multiple Python/Pandas candidates
- Self-Revision - Self-review and correction
- Voting - Multi-candidate voting for best selection
- Code Execution - Safe execution
- Result Verification - Validate result reasonableness

Supports both Chinese and English queries.
"""

import asyncio
import time
import logging
import re
import json
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, field
from enum import Enum

from ..base import BaseEngine, QueryResult

logger = logging.getLogger(__name__)


class TableQAStage(Enum):
    """Table QA pipeline stages"""
    TABLE_LOADING = "table_loading"
    SCHEMA_ANALYSIS = "schema_analysis"
    QUERY_INTERPRETATION = "query_interpretation"
    CODE_GENERATION = "code_generation"
    SELF_REVISION = "self_revision"
    VOTING = "voting"
    CODE_EXECUTION = "code_execution"
    RESULT_VERIFICATION = "result_verification"


@dataclass
class TableQAStep:
    """Single step in Table QA pipeline"""
    stage: TableQAStage
    input: Any
    output: Any
    llm_response: Optional[str] = None
    confidence: float = 0.0
    reasoning: str = ""
    execution_time_ms: float = 0.0
    error: Optional[str] = None


@dataclass
class TableQAResult(QueryResult):
    """Table QA execution result"""
    generated_code: Optional[str] = None
    execution_steps: List[TableQAStep] = field(default_factory=list)
    table_schema: Optional[Dict] = None
    revision_count: int = 0
    voting_results: Optional[Dict] = None
    final_confidence: float = 0.0


@dataclass
class StagePrompt:
    """Prompt template for a generation stage"""
    system: str
    user: str
    examples: List[Dict] = field(default_factory=list)


class MultiStageTableQAEngine(BaseEngine):
    """
    Advanced multi-stage TableQA engine
    
    Complete pipeline:
    1. Table Loading - Load table file
    2. Schema Analysis - Analyze table structure
    3. Query Interpretation - Understand query intent
    4. Code Generation - Parallel generation of multiple code candidates
    5. Self-Revision - Multi-round review and improvement
    6. Voting - Vote to select best code
    7. Code Execution - Safe execution
    8. Result Verification - Validate results
    
    Supports both Chinese and English queries.
    """
    
    # ==================== Prompt Templates (Bilingual) ====================
    
    QUERY_INTERPRETATION_PROMPT = StagePrompt(
        system="You are a data analysis expert. Understand user questions about tables and identify operation types. Support both Chinese and English.",
        user="""Given table schema and a natural language query (Chinese or English), identify the operation the user wants.

Table Schema:
{schema}

Query: {query}

Return JSON:
{{
    "operation": "count|sum|mean|max|min|filter|groupby|sort|select",
    "target_columns": ["col1", "col2"],
    "filters": [{{"column": "col", "condition": ">", "value": 100}}],
    "grouping": ["group_col"],
    "reasoning": "Analysis reasoning"
}}""",
        examples=[],
    )
    
    CODE_GENERATION_PROMPT = StagePrompt(
        system="You are a Python/Pandas expert. Generate code to answer questions about tables. Support both Chinese and English queries. SECURITY: Never output API keys, passwords, or credentials in any form.",
        user="""Given table schema and query interpretation, generate Python/Pandas code.

Table Schema:
{schema}

Query: {query}
Interpretation: {interpretation}

Use 'df' as the DataFrame variable name.
Return ONLY Python code, no explanation.""",
        examples=[
            {
                "query": "What is the total sales? / 销售额总和是多少？",
                "schema": "Columns: sales (float), date (datetime), region (string)",
                "interpretation": "{'operation': 'sum', 'target_columns': ['sales']}",
                "code": "result = df['sales'].sum()",
            },
        ],
    )
    
    REVISION_PROMPT = StagePrompt(
        system="You are a code reviewer. Review generated code and fix issues. Support both Chinese and English.",
        user="""Review this code for correctness.

Query: {query}
Generated Code: {code}
Table Schema: {schema}
Query Interpretation: {interpretation}
{error_message}

If there are issues, write corrected code. If correct, return the original code.

Return ONLY Python code.""",
        examples=[],
    )
    
    VOTING_PROMPT = StagePrompt(
        system="You are a code evaluation expert. Evaluate multiple code candidates and select the best one.",
        user="""Evaluate these code candidates and select the one that best answers the user query.

Query: {query}
Table Schema: {schema}

Candidate Code:
{candidates}

Return JSON:
{{
    "best_index": 0,
    "reasoning": "Reason for selection",
    "scores": [0.9, 0.7, 0.8]
}}""",
        examples=[],
    )
    
    # ==================== Initialization ====================
    
    def __init__(self, config: dict):
        """
        Initialize TableQA engine.
        
        Args:
            config: Configuration containing:
                - table_path: Table file path
                - llm_config: LLM settings
                - generation_config: Generation parameters
        """
        super().__init__(config)
        
        self.table_path = config.get("table_path")
        self.llm_config = config.get("llm_config", {})
        self.generation_config = config.get("generation_config", {
            "num_candidates": 3,
            "max_revisions": 2,
            "parallel_generation": True,
            "voting_enabled": True,
            "safe_execution": True,
        })
        
        self.table_data = None
        self.table_schema = None
        self.llm_client = None
        self.llm_model = self.llm_config.get("model", "gpt-4")
        
        logger.info(f"MultiStageTableQAEngine initialized: {self.table_path}")
    
    def _get_llm_client(self):
        """Lazy load LLM client"""
        if self.llm_client is None:
            try:
                from openai import AsyncOpenAI
                [REDACTED_SECRET]"api_key")
                base_url = self.llm_config.get("base_url")
                if base_url:
                    self.llm_client = AsyncOpenAI([REDACTED_SECRET], base_url=base_url)
                else:
                    self.llm_client = AsyncOpenAI([REDACTED_SECRET])
            except ImportError:
                logger.error("OpenAI not installed")
                return None
        return self.llm_client
    
    def _load_table(self, path: str) -> Any:
        """Load table file"""
        import pandas as pd
        
        if path.endswith('.csv'):
            return pd.read_csv(path, encoding='utf-8')
        elif path.endswith('.xlsx'):
            return pd.read_excel(path)
        elif path.endswith('.tsv'):
            return pd.read_csv(path, sep='\t')
        else:
            raise ValueError(f"Unsupported format: {path}")
    
    def _analyze_schema(self, df: Any) -> Dict:
        """Analyze table schema"""
        import pandas as pd
        
        schema = {
            "columns": [],
            "row_count": len(df),
            "dtypes": {},
            "statistics": {},
        }
        
        for col in df.columns:
            col_info = {
                "name": col,
                "dtype": str(df[col].dtype),
                "non_null": int(df[col].notna().sum()),
                "unique": int(df[col].nunique()),
                "sample_values": df[col].dropna().head(3).tolist() if len(df) > 0 else [],
            }
            schema["columns"].append(col_info)
            schema["dtypes"][col] = str(df[col].dtype)
            
            # Statistics for numeric columns
            if pd.api.types.is_numeric_dtype(df[col]):
                schema["statistics"][col] = {
                    "mean": float(df[col].mean()) if len(df) > 0 else 0,
                    "min": float(df[col].min()) if len(df) > 0 else 0,
                    "max": float(df[col].max()) if len(df) > 0 else 0,
                    "std": float(df[col].std()) if len(df) > 1 else 0,
                }
        
        return schema
    
    async def execute(
        self,
        query: str,
        context: dict,
        model: Optional[str] = None,
        [REDACTED_SECRET] = None,
    ) -> TableQAResult:
        """
        Execute complete TableQA multi-stage pipeline.
        
        Args:
            query: Natural language query
            context: Execution context
            model: Model name to use (overrides config)
            api_key: API key to use (overrides config)
        """
        # Override LLM config if provided
        if model or [REDACTED_SECRET] = self.llm_config.copy()
            if model:
                llm_config["model"] = model
            if [REDACTED_SECRET]"api_key"] = api_key
            self.llm_client = None  # Reset client
        start_time = time.time()
        execution_steps = []
        
        logger.info(f"Executing TableQA: {query}")
        
        try:
            # Step 1: Table Loading
            load_step = await self._stage_table_loading(context)
            execution_steps.append(load_step)
            
            if load_step.output is None:
                raise ValueError("Failed to load table")
            
            # Step 2: Schema Analysis
            schema_step = await self._stage_schema_analysis(load_step.output)
            execution_steps.append(schema_step)
            self.table_schema = schema_step.output
            
            # Step 3: Query Interpretation
            interpretation_step = await self._stage_query_interpretation(
                query, self.table_schema
            )
            execution_steps.append(interpretation_step)
            interpretation = interpretation_step.output
            
            # Step 4: Code Generation (parallel)
            generation_step = await self._stage_code_generation(
                query, interpretation, self.table_schema
            )
            execution_steps.append(generation_step)
            code_candidates = generation_step.output
            
            if not code_candidates:
                raise ValueError("Failed to generate any code candidates")
            
            # Step 5: Self-Revision
            revision_step = await self._stage_self_revision(
                query, code_candidates, self.table_schema, interpretation
            )
            execution_steps.append(revision_step)
            revised_code = revision_step.output
            
            # Step 6: Voting
            best_code = revised_code[0] if revised_code else None
            voting_results = None
            
            if self.generation_config.get("voting_enabled", True) and len(revised_code) > 1:
                voting_step = await self._stage_voting(
                    query, revised_code, self.table_schema, interpretation
                )
                execution_steps.append(voting_step)
                best_code = voting_step.output["best_code"]
                voting_results = voting_step.output
            
            # Step 7: Code Execution
            execution_step = await self._stage_code_execution(
                best_code,
                load_step.output,
            )
            execution_steps.append(execution_step)
            
            # Step 8: Result Verification
            verification_step = await self._stage_result_verification(
                query, execution_step.output, best_code
            )
            execution_steps.append(verification_step)
            
            # Build result
            total_time = (time.time() - start_time) * 1000
            
            result = TableQAResult(
                data=execution_step.output.get("result"),
                success=execution_step.output.get("success", False),
                engine=self.name,
                query=query,
                execution_time_ms=total_time,
                confidence=verification_step.confidence,
                generated_code=best_code,
                execution_steps=execution_steps,
                table_schema=self.table_schema,
                revision_count=len(revision_step.output),
                voting_results=voting_results,
                final_confidence=verification_step.confidence,
                error=execution_step.output.get("error"),
            )
            
            logger.info(f"TableQA complete: {total_time:.0f}ms, "
                       f"confidence: {result.confidence:.2f}")
            
            return result
            
        except Exception as e:
            logger.error(f"TableQA execution failed: {e}")
            import traceback
            traceback.print_exc()
            return TableQAResult(
                data=None,
                success=False,
                engine=self.name,
                query=query,
                execution_time_ms=(time.time() - start_time) * 1000,
                confidence=0.0,
                generated_code=None,
                execution_steps=execution_steps,
                table_schema=None,
                revision_count=0,
                voting_results=None,
                final_confidence=0.0,
                error=str(e),
            )
    
    # ==================== Stage 1: Table Loading ====================
    
    async def _stage_table_loading(self, context: dict) -> TableQAStep:
        """Stage 1: Load table"""
        start = time.time()
        
        try:
            path = context.get("table_path") or self.table_path
            
            if not path:
                return TableQAStep(
                    stage=TableQAStage.TABLE_LOADING,
                    input=context,
                    output=None,
                    confidence=0.0,
                    reasoning="No table path provided",
                    execution_time_ms=(time.time() - start) * 1000,
                    error="No table path",
                )
            
            df = self._load_table(path)
            
            return TableQAStep(
                stage=TableQAStage.TABLE_LOADING,
                input=path,
                output=df,
                confidence=0.95,
                reasoning=f"Loaded {len(df)} rows, {len(df.columns)} columns",
                execution_time_ms=(time.time() - start) * 1000,
            )
            
        except Exception as e:
            return TableQAStep(
                stage=TableQAStage.TABLE_LOADING,
                input=context,
                output=None,
                confidence=0.0,
                reasoning=f"Load failed: {str(e)}",
                execution_time_ms=(time.time() - start) * 1000,
                error=str(e),
            )
    
    # ==================== Stage 2: Schema Analysis ====================
    
    async def _stage_schema_analysis(self, df: Any) -> TableQAStep:
        """Stage 2: Analyze schema"""
        start = time.time()
        
        schema = self._analyze_schema(df)
        
        step = TableQAStep(
            stage=TableQAStage.SCHEMA_ANALYSIS,
            input=f"{len(df)} rows, {len(df.columns)} columns",
            output=schema,
            confidence=0.95,
            reasoning=f"Analyzed {len(schema['columns'])} columns",
            execution_time_ms=(time.time() - start) * 1000,
        )
        
        logger.debug(f"Stage 2 (Schema Analysis): {step.reasoning}")
        
        return step
    
    # ==================== Stage 3: Query Interpretation ====================
    
    async def _stage_query_interpretation(
        self,
        query: str,
        schema: Dict,
    ) -> TableQAStep:
        """Stage 3: Query interpretation"""
        start = time.time()
        
        # Use LLM for query interpretation
        client = self._get_llm_client()
        interpretation = None
        
        if client:
            schema_str = self._format_schema_for_prompt(schema)
            prompt = self.QUERY_INTERPRETATION_PROMPT.user.format(
                schema=schema_str,
                query=query,
            )
            
            try:
                response = await client.chat.completions.create(
                    model=self.llm_model,
                    messages=[
                        {"role": "system", "content": self.QUERY_INTERPRETATION_PROMPT.system},
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.1,
                    max_tokens=500,
                    response_format={"type": "json_object"},
                )
                
                result = response.choices[0].message.content.strip()
                result = self._extract_json(result)
                
                if result:
                    interpretation = result
                
            except Exception as e:
                logger.warning(f"LLM query interpretation failed, using rule-based: {e}")
        
        # Fallback: rule-based interpretation
        if interpretation is None:
            interpretation = {
                "operation": self._detect_operation(query),
                "target_columns": self._detect_columns(query, schema),
                "filters": self._detect_filters(query, schema),
                "grouping": self._detect_grouping(query, schema),
                "reasoning": "Rule-based matching",
            }
        
        step = TableQAStep(
            stage=TableQAStage.QUERY_INTERPRETATION,
            input=query,
            output=interpretation,
            confidence=0.8 if interpretation else 0.5,
            reasoning=f"Detected operation: {interpretation.get('operation', 'unknown')}",
            execution_time_ms=(time.time() - start) * 1000,
        )
        
        logger.debug(f"Stage 3 (Query Interpretation): {step.reasoning}")
        
        return step
    
    def _detect_operation(self, query: str) -> str:
        """Detect operation type from query - supports Chinese and English"""
        query_lower = query.lower()
        
        # Bilingual operation keywords
        operations = [
            ("count", ["多少", "几个", "数量", "count", "how many"]),
            ("sum", ["总和", "合计", "总计", "sum", "total"]),
            ("mean", ["平均", "均值", "mean", "average"]),
            ("max", ["最大", "最高", "max", "maximum"]),
            ("min", ["最小", "最低", "min", "minimum"]),
            ("filter", ["筛选", "过滤", "where", "filter"]),
            ("groupby", ["分组", "按", "group by"]),
            ("sort", ["排序", "顺序", "order by", "sort"]),
        ]
        
        for op, keywords in operations:
            if any(kw in query_lower for kw in keywords):
                return op
        
        return "select"
    
    def _detect_columns(self, query: str, schema: Dict) -> List[str]:
        """Detect relevant columns from query"""
        matched = []
        query_lower = query.lower()
        
        for col in schema["columns"]:
            if col["name"].lower() in query_lower:
                matched.append(col["name"])
        
        return matched
    
    def _detect_filters(self, query: str, schema: Dict) -> List[Dict]:
        """Detect filter conditions"""
        filters = []
        
        # Simple pattern matching
        patterns = [
            (r'(大于 | 超过|greater|more than|>)\s*(\d+(?:\.\d+)?)', ">"),
            (r'(小于 | 低于|less|less than|<)\s*(\d+(?:\.\d+)?)', "<"),
            (r'(等于 | 等于|equals|=)\s*(\d+(?:\.\d+)?)', "="),
        ]
        
        for pattern, op in patterns:
            match = re.search(pattern, query, re.IGNORECASE)
            if match:
                filters.append({
                    "operator": op,
                    "value": match.group(2),
                })
        
        return filters
    
    def _detect_grouping(self, query: str, schema: Dict) -> List[str]:
        """Detect grouping columns"""
        # Simple implementation
        return []
    
    # ==================== Stage 4: Code Generation ====================
    
    async def _stage_code_generation(
        self,
        query: str,
        interpretation: Dict,
        schema: Dict,
    ) -> TableQAStep:
        """Stage 4: Parallel generation of code candidates"""
        start = time.time()
        
        num_candidates = self.generation_config.get("num_candidates", 3)
        parallel = self.generation_config.get("parallel_generation", True)
        
        if parallel:
            tasks = [
                self._generate_single_code(query, interpretation, schema, i)
                for i in range(num_candidates)
            ]
            candidates = await asyncio.gather(*tasks, return_exceptions=True)
            candidates = [c for c in candidates if isinstance(c, str) and c.strip()]
        else:
            candidates = []
            for i in range(num_candidates):
                code = await self._generate_single_code(
                    query, interpretation, schema, i
                )
                if code and code.strip():
                    candidates.append(code)
        
        step = TableQAStep(
            stage=TableQAStage.CODE_GENERATION,
            input={"query": query, "interpretation": interpretation},
            output=candidates,
            confidence=0.7,
            reasoning=f"Generated {len(candidates)} code candidates",
            execution_time_ms=(time.time() - start) * 1000,
        )
        
        logger.debug(f"Stage 4 (Code Generation): {len(candidates)} candidates")
        
        return step
    
    async def _generate_single_code(
        self,
        query: str,
        interpretation: Dict,
        schema: Dict,
        candidate_id: int,
    ) -> Optional[str]:
        """Generate a single code candidate"""
        client = self._get_llm_client()
        if not client:
            return self._generate_code_template(interpretation, schema)
        
        schema_str = self._format_schema_for_prompt(schema)
        
        prompt = self.CODE_GENERATION_PROMPT.user.format(
            schema=schema_str,
            query=query,
            interpretation=interpretation,
        )
        
        try:
            response = await client.chat.completions.create(
                model=self.llm_model,
                messages=[
                    {"role": "system", "content": self.CODE_GENERATION_PROMPT.system},
                    {"role": "user", "content": prompt},
                ],
                temperature=0.3 + (candidate_id * 0.15),
                max_tokens=400,
            )
            
            code = response.choices[0].message.content.strip()
            code = self._extract_code_from_response(code)
            
            return code
            
        except Exception as e:
            logger.error(f"Code generation failed: {e}")
            return self._generate_code_template(interpretation, schema)
    
    def _generate_code_template(
        self,
        interpretation: Dict,
        schema: Dict,
    ) -> str:
        """Generate code template (fallback)"""
        op = interpretation.get("operation", "select")
        cols = interpretation.get("target_columns", [])
        
        if not cols and schema["columns"]:
            cols = [schema["columns"][0]["name"]]
        
        col = cols[0] if cols else "col"
        
        templates = {
            "count": "result = len(df)",
            "sum": f"result = df['{col}'].sum()",
            "mean": f"result = df['{col}'].mean()",
            "max": f"result = df['{col}'].max()",
            "min": f"result = df['{col}'].min()",
            "filter": f"result = df[df['{col}'].notna()]",
            "groupby": f"result = df.groupby('{col}').size()",
            "sort": f"result = df.sort_values('{col}')",
            "select": f"result = df[['{col}']]",
        }
        
        return templates.get(op, "result = df")
    
    # ==================== Stage 5: Self-Revision ====================
    
    async def _stage_self_revision(
        self,
        query: str,
        candidates: List[str],
        schema: Dict,
        interpretation: Dict,
    ) -> TableQAStep:
        """Stage 5: Multi-round self-revision"""
        start = time.time()
        revised = candidates.copy()
        max_revisions = self.generation_config.get("max_revisions", 2)
        
        for round_num in range(max_revisions):
            logger.debug(f"Revision round {round_num + 1}/{max_revisions}")
            
            new_revised = []
            
            for code in revised:
                revised_code = await self._revise_single_code(
                    query, code, schema, interpretation, round_num
                )
                new_revised.append(revised_code)
            
            if new_revised == revised:
                logger.debug(f"Revision converged at round {round_num + 1}")
                break
            
            revised = new_revised
        
        step = TableQAStep(
            stage=TableQAStage.SELF_REVISION,
            input=candidates,
            output=revised,
            confidence=0.75,
            reasoning=f"Completed revision of {len(revised)} candidates",
            execution_time_ms=(time.time() - start) * 1000,
        )
        
        logger.debug(f"Stage 5 (Self-Revision): {len(revised)} revised candidates")
        
        return step
    
    async def _revise_single_code(
        self,
        query: str,
        code: str,
        schema: Dict,
        interpretation: Dict,
        iteration: int,
    ) -> str:
        """Revise a single code candidate"""
        client = self._get_llm_client()
        if not client:
            return code
        
        # Try to validate code
        error_message = ""
        try:
            scope = {"df": None, "pd": __import__("pandas")}
            compile(code, '<string>', 'exec')
        except SyntaxError as e:
            error_message = f"Syntax error: {str(e)}"
        
        schema_str = self._format_schema_for_prompt(schema)
        
        prompt = self.REVISION_PROMPT.user.format(
            query=query,
            code=code,
            schema=schema_str,
            interpretation=interpretation,
            error_message=error_message,
        )
        
        try:
            response = await client.chat.completions.create(
                model=self.llm_model,
                messages=[
                    {"role": "system", "content": self.REVISION_PROMPT.system},
                    {"role": "user", "content": prompt},
                ],
                temperature=0.2,
                max_tokens=400,
            )
            
            revised_code = response.choices[0].message.content.strip()
            revised_code = self._extract_code_from_response(revised_code)
            
            return revised_code
            
        except Exception as e:
            logger.error(f"Revision failed: {e}")
            return code
    
    # ==================== Stage 6: Voting ====================
    
    async def _stage_voting(
        self,
        query: str,
        candidates: List[str],
        schema: Dict,
        interpretation: Dict,
    ) -> TableQAStep:
        """Stage 6: Vote to select best code"""
        start = time.time()
        
        client = self._get_llm_client()
        voting_result = None
        
        if client:
            schema_str = self._format_schema_for_prompt(schema)
            candidates_str = "\n\n".join([f"Candidate {i}:\n{code}" for i, code in enumerate(candidates)])
            
            prompt = self.VOTING_PROMPT.user.format(
                query=query,
                schema=schema_str,
                candidates=candidates_str,
            )
            
            try:
                response = await client.chat.completions.create(
                    model=self.llm_model,
                    messages=[
                        {"role": "system", "content": self.VOTING_PROMPT.system},
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.1,
                    max_tokens=500,
                    response_format={"type": "json_object"},
                )
                
                result = response.choices[0].message.content.strip()
                result = self._extract_json(result)
                
                if result and "best_index" in result:
                    best_idx = result["best_index"]
                    if 0 <= best_idx < len(candidates):
                        voting_result = {
                            "best_index": best_idx,
                            "best_code": candidates[best_idx],
                            "best_score": result.get("scores", [0.5] * len(candidates))[best_idx] if "scores" in result else 0.8,
                            "all_results": [
                                {"index": i, "code": code, "score": result.get("scores", [0.5] * len(candidates))[i] if "scores" in result else 0.5}
                                for i, code in enumerate(candidates)
                            ],
                            "reasoning": result.get("reasoning", ""),
                        }
                
            except Exception as e:
                logger.warning(f"LLM voting failed, using rule-based: {e}")
        
        # Fallback: rule-based voting
        if not voting_result:
            voting_result = {
                "best_index": 0,
                "best_code": candidates[0] if candidates else None,
                "best_score": 0.7,
                "all_results": [{"index": i, "code": code, "score": 0.5} for i, code in enumerate(candidates)],
                "reasoning": "Default to first candidate",
            }
        
        step = TableQAStep(
            stage=TableQAStage.VOTING,
            input=candidates,
            output=voting_result,
            confidence=voting_result["best_score"],
            reasoning=f"Selected candidate #{voting_result['best_index']}",
            execution_time_ms=(time.time() - start) * 1000,
        )
        
        logger.debug(f"Stage 6 (Voting): Selected #{voting_result['best_index']}")
        
        return step
    
    # ==================== Stage 7: Code Execution ====================
    
    async def _stage_code_execution(
        self,
        code: Optional[str],
        df: Any,
    ) -> TableQAStep:
        """Stage 7: Execute code"""
        start = time.time()
        
        if not code:
            return TableQAStep(
                stage=TableQAStage.CODE_EXECUTION,
                input=code,
                output={"success": False, "result": None, "error": "No code"},
                confidence=0.0,
                reasoning="No code to execute",
                execution_time_ms=(time.time() - start) * 1000,
                error="No code generated",
            )
        
        try:
            # Safe execution
            import pandas as pd
            scope = {"df": df, "pd": pd}
            
            exec(code, scope)
            result = scope.get("result")
            
            return TableQAStep(
                stage=TableQAStage.CODE_EXECUTION,
                input=code,
                output={"success": True, "result": result},
                confidence=0.85,
                reasoning="Code executed successfully",
                execution_time_ms=(time.time() - start) * 1000,
            )
            
        except Exception as e:
            return TableQAStep(
                stage=TableQAStage.CODE_EXECUTION,
                input=code,
                output={"success": False, "result": None, "error": str(e)},
                confidence=0.2,
                reasoning=f"Execution failed: {str(e)}",
                execution_time_ms=(time.time() - start) * 1000,
                error=str(e),
            )
    
    # ==================== Stage 8: Result Verification ====================
    
    async def _stage_result_verification(
        self,
        query: str,
        execution_result: Dict,
        code: Optional[str],
    ) -> TableQAStep:
        """Stage 8: Result verification"""
        start = time.time()
        
        confidence = 0.0
        
        if execution_result.get("success"):
            result = execution_result.get("result")
            if result is not None:
                confidence = 0.85
            else:
                confidence = 0.5
        else:
            confidence = 0.2
        
        step = TableQAStep(
            stage=TableQAStage.RESULT_VERIFICATION,
            input=execution_result,
            output={"confidence": confidence},
            confidence=confidence,
            reasoning=f"Verification: confidence={confidence:.2f}",
            execution_time_ms=(time.time() - start) * 1000,
        )
        
        logger.debug(f"Stage 8 (Result Verification): confidence={confidence:.2f}")
        
        return step
    
    # ==================== Utilities ====================
    
    def _format_schema_for_prompt(self, schema: Dict) -> str:
        """Format schema for prompt"""
        lines = []
        
        for col in schema["columns"]:
            dtype = col["dtype"]
            non_null = col["non_null"]
            lines.append(f"Column: {col['name']} ({dtype}, {non_null} non-null)")
        
        lines.append(f"Total rows: {schema['row_count']}")
        
        return "\n".join(lines)
    
    def _extract_code_from_response(self, response: str) -> str:
        """Extract code from response"""
        code = response.strip()
        
        if "```python" in code:
            code = code.split("```python")[1].split("```")[0].strip()
        elif "```" in code:
            code = code.split("```")[1].split("```")[0].strip()
        
        return code
    
    def _extract_json(self, text: str) -> Optional[Dict]:
        """Extract JSON from text"""
        import json
        import re
        
        try:
            return json.loads(text)
        except:
            pass
        
        match = re.search(r'\{.*\}', text, re.DOTALL)
        if match:
            try:
                return json.loads(match.group())
            except:
                pass
        
        return None
    
    def get_confidence(self, query: str, context: dict) -> float:
        """Get confidence for this query"""
        if len(query.split()) > 3:
            return 0.7
        return 0.5


path: src/engines/nl2sparql/__init__.py
content: """
HeteroMind - NL2SPARQL Engine Package

NL2SPARQL engine implementations.
"""

from .base_engine import BaseNL2SPARQLEngine, NL2SPARQLResult
from .multi_stage_engine import MultiStageNL2SPARQLEngine

__all__ = [
    "BaseNL2SPARQLEngine",
    "NL2SPARQLResult",
    "MultiStageNL2SPARQLEngine",
]


path: src/engines/nl2sparql/base_engine.py
content: """
HeteroMind - NL2SPARQL Base Engine

Base class and result types for NL2SPARQL engines.
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from enum import Enum
import logging

from ..base import BaseEngine, QueryResult

logger = logging.getLogger(__name__)


class SPARQLGenerationStage(Enum):
    """Stages in SPARQL generation pipeline"""
    ENTITY_LINKING = "entity_linking"
    ONTOLOGY_RETRIEVAL = "ontology_retrieval"
    INITIAL_GENERATION = "initial_generation"
    SELF_REVISION = "self_revision"
    VALIDATION = "validation"
    EXECUTION = "execution"
    RESULT_VERIFICATION = "result_verification"


@dataclass
class SPARQLGenerationStep:
    """Single step in SPARQL generation"""
    stage: SPARQLGenerationStage
    input: Any
    output: Any
    llm_response: Optional[str] = None
    confidence: float = 0.0
    reasoning: str = ""
    execution_time_ms: float = 0.0
    error: Optional[str] = None


@dataclass
class LinkedEntity:
    """Entity linked to knowledge graph"""
    mention: str
    uri: str
    label: str
    score: float
    types: List[str] = field(default_factory=list)


@dataclass
class NL2SPARQLResult(QueryResult):
    """Extended QueryResult for NL2SPARQL"""
    generated_sparql: Optional[str] = None
    execution_steps: List[SPARQLGenerationStep] = field(default_factory=list)
    linked_entities: List[LinkedEntity] = field(default_factory=list)
    ontology_used: Optional[Dict] = None
    revision_count: int = 0
    voting_results: Optional[List[Dict]] = None
    final_confidence: float = 0.0


class BaseNL2SPARQLEngine(BaseEngine):
    """
    Base class for NL2SPARQL engines.
    
    Provides common utilities for:
    - Entity linking to KG
    - Ontology retrieval
    - Multi-stage SPARQL generation
    - Self-revision mechanisms
    - Voting between candidates
    - Result validation
    """
    
    def __init__(self, config: dict):
        """
        Initialize NL2SPARQL engine.
        
        Args:
            config: Configuration dict with:
                - endpoint_url: SPARQL endpoint URL
                - ontology: KG ontology metadata
                - llm_config: LLM configuration
                - generation_config: Generation parameters
        """
        super().__init__(config)
        
        self.endpoint_url = config.get("endpoint_url")
        self.ontology = config.get("ontology", {})
        self.llm_config = config.get("llm_config", {})
        self.generation_config = config.get("generation_config", {
            "num_candidates": 3,
            "max_revisions": 2,
            "voting_enabled": True,
            "parallel_generation": True,
            "entity_linking_threshold": 0.7,
        })
        
        # Ontology index for efficient retrieval
        self.ontology_index = self._build_ontology_index()
    
    def _build_ontology_index(self) -> Dict:
        """Build searchable ontology index"""
        index = {
            "classes": {},
            "properties": {},
            "instances": {},
            "relationships": [],
        }
        
        # Index classes
        for cls in self.ontology.get("classes", []):
            class_name = cls["name"].lower()
            index["classes"][class_name] = cls
            
            # Index labels and synonyms
            for label in cls.get("labels", []):
                index["classes"][label.lower()] = cls
            
            for synonym in cls.get("synonyms", []):
                index["classes"][synonym.lower()] = cls
        
        # Index properties
        for prop in self.ontology.get("properties", []):
            prop_name = prop["name"].lower()
            index["properties"][prop_name] = prop
            
            # Index labels
            for label in prop.get("labels", []):
                index["properties"][label.lower()] = prop
        
        # Index relationships (object properties)
        for rel in self.ontology.get("relationships", []):
            index["relationships"].append(rel)
            
            # Index by predicate name
            predicate = rel.get("predicate", "").lower()
            if predicate:
                index["properties"][predicate] = rel
        
        logger.debug(f"Built ontology index: {len(index['classes'])} classes, "
                    f"{len(index['properties'])} properties")
        
        return index
    
    async def link_entities(self, query: str) -> List[LinkedEntity]:
        """
        Link entity mentions in query to knowledge graph.
        
        Args:
            query: Natural language query
            
        Returns:
            List of linked entities
        """
        # Extract entity mentions (simple NER)
        mentions = self._extract_entity_mentions(query)
        
        linked = []
        threshold = self.generation_config.get("entity_linking_threshold", 0.7)
        
        for mention in mentions:
            # Search ontology for match
            best_match = self._find_best_entity_match(mention)
            
            if best_match and best_match["score"] >= threshold:
                linked.append(LinkedEntity(
                    mention=mention,
                    uri=best_match["uri"],
                    label=best_match["label"],
                    score=best_match["score"],
                    types=best_match.get("types", []),
                ))
        
        logger.debug(f"Linked {len(linked)}/{len(mentions)} entities")
        
        return linked
    
    def _extract_entity_mentions(self, query: str) -> List[str]:
        """Extract potential entity mentions from query"""
        import re
        
        mentions = []
        
        # Quoted strings
        quoted = re.findall(r'"([^"]+)"|\'([^\']+)\'', query)
        mentions.extend([q[0] or q[1] for q in quoted if q[0] or q[1]])
        
        # Capitalized phrases (simple NER)
        capitalized = re.findall(r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b', query)
        mentions.extend(capitalized)
        
        # Remove duplicates
        seen = set()
        unique = []
        for m in mentions:
            if m not in seen and len(m) > 1:
                seen.add(m)
                unique.append(m)
        
        return unique
    
    def _find_best_entity_match(self, mention: str) -> Optional[Dict]:
        """Find best ontology match for entity mention"""
        mention_lower = mention.lower()
        best_match = None
        best_score = 0.0
        
        # Check classes
        for class_name, cls in self.ontology_index["classes"].items():
            if mention_lower == class_name:
                return {
                    "uri": cls.get("uri", f"http://example.org/{class_name}"),
                    "label": cls.get("label", mention),
                    "score": 1.0,
                    "types": [cls.get("name")],
                }
            elif mention_lower in class_name or class_name in mention_lower:
                score = 0.8
                if score > best_score:
                    best_score = score
                    best_match = {
                        "uri": cls.get("uri"),
                        "label": cls.get("label", mention),
                        "score": score,
                        "types": [cls.get("name")],
                    }
        
        # Check properties
        for prop_name, prop in self.ontology_index["properties"].items():
            if mention_lower == prop_name:
                return {
                    "uri": prop.get("uri"),
                    "label": prop.get("label", mention),
                    "score": 1.0,
                    "types": ["Property"],
                }
        
        return best_match if best_score >= 0.5 else None
    
    def retrieve_relevant_ontology(
        self,
        query: str,
        linked_entities: List[LinkedEntity],
    ) -> Dict:
        """
        Retrieve ontology elements relevant to the query.
        
        Args:
            query: Natural language query
            linked_entities: Linked entity mentions
            
        Returns:
            Relevant ontology subset
        """
        relevant = {
            "classes": [],
            "properties": [],
            "relationships": [],
        }
        
        # Add types of linked entities
        for entity in linked_entities:
            for type_uri in entity.types:
                cls = self.ontology_index["classes"].get(type_uri.lower())
                if cls and cls not in relevant["classes"]:
                    relevant["classes"].append(cls)
        
        # Find properties based on query keywords
        query_lower = query.lower()
        for prop_name, prop in self.ontology_index["properties"].items():
            if prop_name in query_lower:
                relevant["properties"].append(prop)
        
        # Find relationships for matched classes
        class_uris = {c.get("uri") for c in relevant["classes"]}
        for rel in self.ontology_index["relationships"]:
            domain = rel.get("domain", "")
            range_ = rel.get("range", "")
            if domain in class_uris or range_ in class_uris:
                relevant["relationships"].append(rel)
        
        logger.debug(f"Retrieved ontology: {len(relevant['classes'])} classes, "
                    f"{len(relevant['properties'])} properties")
        
        return relevant
    
    async def generate_sparql_candidates(
        self,
        query: str,
        ontology: Dict,
        linked_entities: List[LinkedEntity],
        num_candidates: int = 3,
    ) -> List[str]:
        """
        Generate multiple SPARQL candidates.
        
        Args:
            query: Natural language query
            ontology: Relevant ontology subset
            linked_entities: Linked entities
            num_candidates: Number of candidates
            
        Returns:
            List of SPARQL query strings
        """
        raise NotImplementedError
    
    async def revise_sparql(
        self,
        query: str,
        sparql: str,
        ontology: Dict,
        error: Optional[str] = None,
    ) -> str:
        """
        Revise SPARQL based on feedback.
        
        Args:
            query: Original query
            sparql: Current SPARQL
            ontology: Ontology subset
            error: Error message if any
            
        Returns:
            Revised SPARQL
        """
        raise NotImplementedError
    
    def validate_sparql(self, sparql: str) -> Dict:
        """
        Validate SPARQL syntax and semantics.
        
        Args:
            sparql: SPARQL query string
            
        Returns:
            Validation result with errors/warnings
        """
        validation = {
            "valid": True,
            "errors": [],
            "warnings": [],
        }
        
        # Basic syntax checks
        sparql_upper = sparql.upper().strip()
        
        # Must contain SELECT or ASK
        if "SELECT" not in sparql_upper and "ASK" not in sparql_upper:
            validation["valid"] = False
            validation["errors"].append("Missing SELECT or ASK clause")
        
        # Must contain WHERE
        if "WHERE" not in sparql_upper:
            validation["valid"] = False
            validation["errors"].append("Missing WHERE clause")
        
        # Check for balanced braces
        if sparql.count("{") != sparql.count("}"):
            validation["valid"] = False
            validation["errors"].append("Unbalanced braces")
        
        # Check for PREFIX declarations (warning only)
        if "PREFIX" not in sparql_upper and "http://" in sparql.lower():
            validation["warnings"].append("Consider adding PREFIX declarations")
        
        return validation
    
    async def vote_on_candidates(
        self,
        query: str,
        candidates: List[str],
        ontology: Dict,
        linked_entities: List[LinkedEntity],
    ) -> Dict:
        """
        Vote on multiple SPARQL candidates.
        
        Args:
            query: Original query
            candidates: SPARQL candidates
            ontology: Ontology subset
            linked_entities: Linked entities
            
        Returns:
            Voting results
        """
        voting_results = []
        
        for i, sparql in enumerate(candidates):
            score = self._score_candidate(
                query, sparql, ontology, linked_entities
            )
            voting_results.append({
                "index": i,
                "sparql": sparql,
                "score": score,
                "criteria": self._get_scoring_criteria(),
            })
        
        best = max(voting_results, key=lambda x: x["score"])
        
        return {
            "best_index": best["index"],
            "best_sparql": best["sparql"],
            "best_score": best["score"],
            "all_results": voting_results,
        }
    
    def _score_candidate(
        self,
        query: str,
        sparql: str,
        ontology: Dict,
        linked_entities: List[LinkedEntity],
    ) -> float:
        """Score a SPARQL candidate"""
        score = 0.0
        sparql_lower = sparql.lower()
        
        # Criterion 1: Uses linked entities
        for entity in linked_entities:
            if entity.uri.lower() in sparql_lower or \
               entity.label.lower() in sparql_lower:
                score += 0.2
        
        # Criterion 2: Has proper structure
        if "SELECT" in sparql.upper() and "WHERE" in sparql.upper():
            score += 0.3
        
        # Criterion 3: Valid syntax
        validation = self.validate_sparql(sparql)
        if validation["valid"]:
            score += 0.3
        else:
            score -= len(validation["errors"]) * 0.2
        
        # Criterion 4: Uses ontology properties
        for prop in ontology.get("properties", []):
            if prop.get("name", "").lower() in sparql_lower:
                score += 0.1
        
        return min(1.0, max(0.0, score))
    
    def _get_scoring_criteria(self) -> Dict:
        """Get scoring criteria descriptions"""
        return {
            "entity_use": "Uses linked entities",
            "structure": "Has proper SPARQL structure",
            "validity": "No syntax errors",
            "ontology_use": "Uses ontology properties",
        }


path: src/engines/__init__.py
content: """
HeteroMind - Engines Package

Knowledge source engine implementations.
"""

from .base import BaseEngine, BaseSQLEngine, BaseSPARQLEngine, BaseTableQAEngine, QueryResult
from .nl2sql import BaseNL2SQLEngine, NL2SQLResult, MultiStageNL2SQLEngine
from .nl2sparql import BaseNL2SPARQLEngine, NL2SPARQLResult, MultiStageNL2SPARQLEngine
from .table_qa import BaseTableQAEngine, TableQAResult, MultiStageTableQAEngine

__all__ = [
    # Base classes
    "BaseEngine",
    "BaseSQLEngine",
    "BaseSPARQLEngine",
    "BaseTableQAEngine",
    "QueryResult",
    # NL2SQL
    "BaseNL2SQLEngine",
    "NL2SQLResult",
    "MultiStageNL2SQLEngine",
    # NL2SPARQL
    "BaseNL2SPARQLEngine",
    "NL2SPARQLResult",
    "MultiStageNL2SPARQLEngine",
    # Table QA
    "BaseTableQAEngine",
    "TableQAResult",
    "MultiStageTableQAEngine",
]


path: src/engines/nl2sql/__init__.py
content: """
HeteroMind - NL2SQL Engine Package

NL2SQL engine implementations.
"""

from .base_engine import BaseNL2SQLEngine, NL2SQLResult
from .multi_stage_engine import MultiStageNL2SQLEngine

__all__ = [
    "BaseNL2SQLEngine",
    "NL2SQLResult",
    "MultiStageNL2SQLEngine",
]


path: src/engines/nl2sql/base_engine.py
content: """
HeteroMind - NL2SQL Base Engine

Base class and result types for NL2SQL engines.
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from enum import Enum
import logging

from ..base import BaseEngine, QueryResult

logger = logging.getLogger(__name__)


class SQLGenerationStage(Enum):
    """Stages in SQL generation pipeline"""
    SCHEMA_RETRIEVAL = "schema_retrieval"
    INITIAL_GENERATION = "initial_generation"
    SELF_REVISION = "self_revision"
    VALIDATION = "validation"
    EXECUTION = "execution"
    RESULT_VERIFICATION = "result_verification"


@dataclass
class SQLGenerationStep:
    """Single step in SQL generation"""
    stage: SQLGenerationStage
    input: Any
    output: Any
    llm_response: Optional[str] = None
    confidence: float = 0.0
    reasoning: str = ""
    execution_time_ms: float = 0.0
    error: Optional[str] = None


@dataclass
class NL2SQLResult(QueryResult):
    """Extended QueryResult for NL2SQL"""
    generated_sql: Optional[str] = None
    execution_steps: List[SQLGenerationStep] = field(default_factory=list)
    schema_used: Optional[Dict] = None
    revision_count: int = 0
    voting_results: Optional[List[Dict]] = None
    final_confidence: float = 0.0


class BaseNL2SQLEngine(BaseEngine):
    """
    Base class for NL2SQL engines.
    
    Provides common utilities for:
    - Schema retrieval and indexing
    - Multi-stage SQL generation
    - Self-revision mechanisms
    - Voting between multiple candidates
    - Result validation
    """
    
    def __init__(self, config: dict):
        """
        Initialize NL2SQL engine.
        
        Args:
            config: Configuration dict with:
                - connection_string: Database connection
                - schema: Database schema metadata
                - llm_config: LLM configuration
                - generation_config: Generation parameters
        """
        super().__init__(config)
        
        self.connection_string = config.get("connection_string")
        self.schema = config.get("schema", {})
        self.llm_config = config.get("llm_config", {})
        self.generation_config = config.get("generation_config", {
            "num_candidates": 3,
            "max_revisions": 2,
            "voting_enabled": True,
            "parallel_generation": True,
        })
        
        # Schema index for efficient retrieval
        self.schema_index = self._build_schema_index()
    
    def _build_schema_index(self) -> Dict:
        """Build searchable schema index"""
        index = {
            "tables": {},
            "columns": {},
            "relationships": [],
            "keywords": {},
        }
        
        for table in self.schema.get("tables", []):
            table_name = table["name"].lower()
            index["tables"][table_name] = table
            
            for col in table.get("columns", []):
                col_name = col["name"].lower()
                index["columns"][col_name] = {
                    "table": table_name,
                    "type": col.get("type", "UNKNOWN"),
                    "description": col.get("description", ""),
                }
                
                # Index keywords from column description
                if col.get("description"):
                    for keyword in col["description"].lower().split():
                        if len(keyword) > 3:
                            index["keywords"].setdefault(keyword, []).append({
                                "table": table_name,
                                "column": col_name,
                            })
            
            # Index relationships (foreign keys)
            for fk in table.get("foreign_keys", []):
                index["relationships"].append({
                    "from_table": table_name,
                    "from_column": fk["column"],
                    "to_table": fk["references"]["table"],
                    "to_column": fk["references"]["column"],
                })
        
        logger.debug(f"Built schema index: {len(index['tables'])} tables, "
                    f"{len(index['columns'])} columns")
        
        return index
    
    def retrieve_relevant_schema(self, query: str) -> Dict:
        """
        Retrieve schema elements relevant to the query.
        
        Uses keyword matching and semantic similarity to find
        relevant tables and columns.
        
        Args:
            query: Natural language query
            
        Returns:
            Relevant schema subset
        """
        query_lower = query.lower()
        relevant = {
            "tables": [],
            "columns": [],
            "relationships": [],
        }
        
        # Match table names
        for table_name, table in self.schema_index["tables"].items():
            if table_name in query_lower:
                relevant["tables"].append(table)
                if table.get("description") and table["description"].lower() in query_lower:
                    relevant["tables"].append(table)
        
        # Match column names
        for col_name, col_info in self.schema_index["columns"].items():
            if col_name in query_lower:
                relevant["columns"].append({
                    "name": col_name,
                    **col_info,
                })
        
        # Match keywords from descriptions
        for keyword, matches in self.schema_index["keywords"].items():
            if keyword in query_lower:
                for match in matches:
                    table = self.schema_index["tables"].get(match["table"])
                    if table and table not in relevant["tables"]:
                        relevant["tables"].append(table)
        
        # Find relationships for matched tables
        table_names = {t["name"] for t in relevant["tables"]}
        for rel in self.schema_index["relationships"]:
            if rel["from_table"] in table_names or rel["to_table"] in table_names:
                relevant["relationships"].append(rel)
        
        logger.debug(f"Retrieved schema: {len(relevant['tables'])} tables, "
                    f"{len(relevant['columns'])} columns")
        
        return relevant
    
    async def generate_sql_candidates(
        self,
        query: str,
        schema: Dict,
        num_candidates: int = 3,
    ) -> List[str]:
        """
        Generate multiple SQL candidates.
        
        Args:
            query: Natural language query
            schema: Relevant schema subset
            num_candidates: Number of candidates to generate
            
        Returns:
            List of SQL query strings
        """
        # To be implemented by subclasses
        raise NotImplementedError
    
    async def revise_sql(
        self,
        query: str,
        sql: str,
        schema: Dict,
        error: Optional[str] = None,
    ) -> str:
        """
        Revise SQL based on feedback.
        
        Args:
            query: Original query
            sql: Current SQL
            schema: Schema subset
            error: Error message if any
            
        Returns:
            Revised SQL
        """
        # To be implemented by subclasses
        raise NotImplementedError
    
    def validate_sql(self, sql: str) -> Dict:
        """
        Validate SQL syntax and semantics.
        
        Args:
            sql: SQL query string
            
        Returns:
            Validation result with errors/warnings
        """
        import sqlparse
        
        validation = {
            "valid": True,
            "errors": [],
            "warnings": [],
        }
        
        # Parse SQL
        try:
            parsed = sqlparse.parse(sql)
            if not parsed:
                validation["valid"] = False
                validation["errors"].append("Failed to parse SQL")
                return validation
        except Exception as e:
            validation["valid"] = False
            validation["errors"].append(f"Parse error: {str(e)}")
            return validation
        
        # Check for dangerous operations
        sql_upper = sql.upper()
        dangerous = ["DROP", "DELETE", "TRUNCATE"]
        for keyword in dangerous:
            if keyword in sql_upper:
                validation["warnings"].append(
                    f"Query contains {keyword} - read-only mode enforced"
                )
        
        # Check for SELECT (read queries only)
        if not sql_upper.strip().startswith("SELECT"):
            validation["warnings"].append("Non-SELECT query detected")
        
        return validation
    
    async def vote_on_candidates(
        self,
        query: str,
        candidates: List[str],
        schema: Dict,
    ) -> Dict:
        """
        Vote on multiple SQL candidates to select the best one.
        
        Args:
            query: Original query
            candidates: List of SQL candidates
            schema: Schema subset
            
        Returns:
            Voting results with scores and best candidate
        """
        voting_results = []
        
        for i, sql in enumerate(candidates):
            # Score each candidate
            score = self._score_candidate(query, sql, schema)
            voting_results.append({
                "index": i,
                "sql": sql,
                "score": score,
                "criteria": self._get_scoring_criteria(query, sql, schema),
            })
        
        # Select best candidate
        best = max(voting_results, key=lambda x: x["score"])
        
        logger.debug(f"Voting complete: best candidate #{best['index']} "
                    f"with score {best['score']:.2f}")
        
        return {
            "best_index": best["index"],
            "best_sql": best["sql"],
            "best_score": best["score"],
            "all_results": voting_results,
        }
    
    def _score_candidate(
        self,
        query: str,
        sql: str,
        schema: Dict,
    ) -> float:
        """Score a SQL candidate"""
        score = 0.0
        
        # Criterion 1: Uses relevant tables
        table_names = {t["name"].lower() for t in schema.get("tables", [])}
        sql_lower = sql.lower()
        matched_tables = sum(1 for t in table_names if t in sql_lower)
        score += matched_tables * 0.2
        
        # Criterion 2: Has proper structure (SELECT ... FROM ... WHERE)
        if "SELECT" in sql_upper and "FROM" in sql_upper:
            score += 0.3
        
        # Criterion 3: No syntax errors
        validation = self.validate_sql(sql)
        if validation["valid"]:
            score += 0.3
        else:
            score -= len(validation["errors"]) * 0.2
        
        # Criterion 4: Matches query intent (simple keyword overlap)
        query_keywords = set(query.lower().split())
        sql_keywords = set(sql_lower.split())
        overlap = len(query_keywords & sql_keywords)
        score += min(0.2, overlap * 0.02)
        
        return min(1.0, max(0.0, score))
    
    def _get_scoring_criteria(
        self,
        query: str,
        sql: str,
        schema: Dict,
    ) -> Dict:
        """Get detailed scoring criteria"""
        return {
            "table_match": "Uses relevant tables",
            "structure": "Has proper SQL structure",
            "validity": "No syntax errors",
            "intent_match": "Matches query keywords",
        }


Security scan verdicts — ClawScan: suspicious | static analysis: suspicious | VirusTotal: suspicious | SkillSpector: suspicious (CRITICAL)