expert_system.py

"""
Expert System for multi-subject problem solving with specialized agents.
"""
import torch
import torch.nn as nn
import json
import os
import re
from typing import Dict, List, Tuple, Optional, Set
from pathlib import Path
from transformer import Transformer, create_transformer_model, initialize_weights
from together_ai import TogetherAI, create_together_ai_client
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ExpertSystem:
    """
    Multi-level expert system that routes problems to specialized agents based on subject.
    """
    
    # Default model to use if no specific model is specified for a subject
    DEFAULT_MODEL = "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
    
    # Subject-specific configurations including model names
    SUBJECTS = {
        'general': {
            'description': 'General knowledge and miscellaneous queries',
            'temperature': 0.3,
            'model': "Qwen/Qwen3-32B-FP8",
            'expert_prompt': "You are a knowledgeable assistant with broad general knowledge."
        },
        'physics': {
            'description': 'Physics and related scientific disciplines',
            'temperature': 0.2,
            'model': "deepseek-ai/DeepSeek-R1-Distill-Llama-70B-free",
            'expert_prompt': "You are a physics expert specializing in classical mechanics, quantum physics, and theoretical physics."
        },
        'chemistry': {
            'description': 'Chemistry and chemical processes',
            'temperature': 0.2,
            'model': "Qwen/Qwen3-235B-A22B-Instruct-2507-tput",
            'expert_prompt': "You are a chemistry expert with knowledge of organic, inorganic, physical, and analytical chemistry."
        },
        'maths': {
            'description': 'Mathematics and related fields',
            'temperature': 0.1,  # Lower temperature for precise mathematical reasoning
            'model': "deepseek-ai/DeepSeek-V3",
            'expert_prompt': "You are a mathematics expert specializing in algebra, calculus, and number theory."
        },
        'coding': {
            'description': 'Programming and software development',
            'temperature': 0.3,
            'model': "deepseek-ai/DeepSeek-Coder-33B-instruct",
            'expert_prompt': "You are a senior software engineer with expertise in multiple programming languages and best practices."
        },
        'biology': {
            'description': 'Biological sciences',
            'temperature': 0.2,
            'model': "Qwen/Qwen3-32B-FP8",
            'expert_prompt': "You are a biology expert with knowledge of genetics, evolution, and cellular biology."
        },
        'history': {
            'description': 'Historical events and analysis',
            'temperature': 0.3,
            'model': "Qwen/Qwen3-32B-FP8",
            'expert_prompt': "You are a historian with expertise in world history and historical analysis."
        },
        'literature': {
            'description': 'Literary analysis and writing',
            'temperature': 0.4,  # Slightly more creative for literature
            'model': "Qwen/Qwen3-32B-FP8",
            'expert_prompt': "You are a literature professor with expertise in literary analysis and creative writing."
        },
        'philosophy': {
            'description': 'Philosophical concepts and reasoning',
            'temperature': 0.3,
            'model': "Qwen/Qwen3-32B-FP8",
            'expert_prompt': "You are a philosopher with expertise in both western and eastern philosophical traditions."
        },
        'economics': {
            'description': 'Economic theory and analysis',
            'temperature': 0.2,
            'model': "Qwen/Qwen3-32B-FP8",
            'expert_prompt': "You are an economics expert with knowledge of microeconomics, macroeconomics, and econometrics."
        }
    }
    
    def __init__(
        self, 
        together_ai_key: str, 
        model_name: str = None,
        subject_models: Dict[str, str] = None
    ):
        """
        Initialize the expert system with API key and model configuration.
        
        Args:
            together_ai_key: Together AI API key
            model_name: Default model to use (overrides DEFAULT_MODEL if provided)
            subject_models: Optional dict mapping subjects to specific model names
        """
        # Initialize TogetherAI client
        self.together_ai = create_together_ai_client(api_key=together_ai_key)
        
        # Set default model
        self.default_model = model_name or self.DEFAULT_MODEL
        
        # Update subject models with any custom mappings
        if subject_models:
            for subject, model in subject_models.items():
                if subject in self.SUBJECTS:
                    self.SUBJECTS[subject]['model'] = model
                    logger.info(f"Using custom model '{model}' for subject '{subject}'")
                else:
                    logger.warning(f"Subject '{subject}' not found in SUBJECTS configuration")
    
    def get_subject_config(self, subject: str) -> dict:
        """
        Get configuration for a specific subject.
        
        Args:
            subject: The subject to get configuration for
            
        Returns:
            dict: Configuration for the subject, or default if not found
        """
        subject_lower = subject.lower()
        
        # Try exact match first
        if subject_lower in self.SUBJECTS:
            return self.SUBJECTS[subject_lower]
            
        # Try partial match
        for subj, config in self.SUBJECTS.items():
            if subject_lower in subj or subj in subject_lower:
                return config
                
        # Default to general if no match found
        logger.warning(f"Subject '{subject}' not found, using 'general' configuration")
        return self.SUBJECTS['general']
    
    def generate_response(
        self, 
        prompt: str, 
        subject: str = 'general',
        temperature: float = None,
        max_tokens: int = 1000
    ) -> str:
        """
        Generate a response for the given prompt and subject.
        
        Args:
            prompt: The input prompt
            subject: The subject area (default: 'general')
            temperature: Sampling temperature (optional, overrides subject default if provided)
            max_tokens: Maximum number of tokens to generate
            
        Returns:
            str: Generated response
        """
        # Get subject configuration
        subject_config = self.get_subject_config(subject)
        
        # Use provided temperature or subject default
        temp = temperature if temperature is not None else subject_config.get('temperature', 0.3)
        
        # Get the appropriate model for this subject
        model = subject_config.get('model', self.default_model)
        
        # Add expert prompt if available
        expert_prompt = subject_config.get('expert_prompt', '')
        if expert_prompt:
            prompt = f"{expert_prompt}\n\n{prompt}"
        
        try:
            # Use the TogetherAI wrapper to generate the response
            response = self.together_ai.generate_text(
                prompt=prompt,
                model=model,
                temperature=temp,
                max_tokens=max_tokens
            )
            
            return response
            
        except Exception as e:
            logger.error(f"Error generating response: {str(e)}")
            return f"Error: {str(e)}"


def create_expert_system(
    together_ai_key: str, 
    model_name: str = None,
    subject_models: Dict[str, str] = None
) -> ExpertSystem:
    """
    Factory function to create an ExpertSystem instance with custom model configurations.
    
    Args:
        together_ai_key: Together AI API key
        model_name: Default model to use (overrides DEFAULT_MODEL if provided)
        subject_models: Optional dict mapping subjects to specific model names
            Example: {'maths': 'meta-llama/Meta-Llama-3-70B-Instruct', 'coding': 'deepseek-ai/DeepSeek-Coder-33B-instruct'}
            
    Returns:
        ExpertSystem instance
    """
    return ExpertSystem(
        together_ai_key=together_ai_key,
        model_name=model_name,
        subject_models=subject_models
    )