langchain_deepseek_enhanced.py

import os
os.environ["GTIFF_SRS_SOURCE"]="EPSG"
os.environ['GDAL_DATA'] = '/root/miniconda3/envs/earthagent/share/gdal'
os.environ["PROJ_LIB"]="/root/miniconda3/envs/earthagent/lib/python3.10/site-packages/pyproj/proj_dir/share/proj"


import json
import logging
import asyncio
import argparse
from enum import auto
from tqdm import tqdm
from pathlib import Path
from copy import deepcopy
from datetime import datetime
from logging.handlers import RotatingFileHandler
import time

from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage

# Pprint for debugging
from pprint import pprint

# Change to current directory
os.chdir(os.path.dirname(os.path.abspath(__file__)))

# Global variables
logger = None
temp_dir_path = None

# Configuration - will be loaded from enhanced config
model_name = 'deepseek'
autoplanning = True
num_experiences = 0   # Number of experiences loaded
sys_prompt = None     # Will be loaded from enhanced config


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(
        description="Evaluate Enhanced Earth-Agent with learned experiences"
    )
    parser.add_argument(
        '--enhanced_config',
        type=str,
        required=True,
        help='Path to enhanced config JSON with learned experiences'
    )
    parser.add_argument(
        '--batch_dir',
        type=str,
        default=None,
        help='Batch directory for storing results (e.g., batch_001)'
    )
    parser.add_argument(
        '--question_ids',
        type=str,
        nargs='+',
        default=None,
        help='Specific question IDs to evaluate in this batch'
    )
    return parser.parse_args()


def load_enhanced_config(enhanced_config_path):
    """Load enhanced configuration with learned experiences"""
    global sys_prompt, num_experiences

    print(f"\n{'='*80}")
    print("Loading Enhanced Configuration")
    print(f"{'='*80}")

    with open(enhanced_config_path, 'r') as f:
        enhanced_config = json.load(f)

    # Extract system prompt with experiences
    sys_prompt = enhanced_config['system_prompt']
    num_experiences = enhanced_config['metadata']['num_experiences']

    print(f"Experiment ID: {enhanced_config['exp_id']}")
    print(f"Number of experiences loaded: {num_experiences}")
    print(f"Generated at: {enhanced_config['metadata']['generated_at']}")
    print(f"{'='*80}\n")

    return enhanced_config


def init_global_params(batch_dir=None):
    """Initialize global parameters and logging"""
    global temp_dir_path, logger

    if temp_dir_path is None:
        if batch_dir:
            # Use provided batch directory
            temp_dir_path = Path(batch_dir).absolute()
        else:
            # Use default directory structure
            temp_dir_path = Path('./evaluate_langchain/{}_AP_enhanced_{}exp_{}'.format(
                model_name,
                num_experiences,
                datetime.now().strftime('%y-%m-%d_%H-%M')
            )).absolute()
    temp_dir_path.mkdir(parents=True, exist_ok=True)

    class JsonFormatter(logging.Formatter):
        def format(self, record):
            # Simplified logging compatible with original format
            log_record = {
                "question_index": record.args[0] if record.args else "unknown",
                "timestamp": self.formatTime(record, self.datefmt),
                "conversations": record.args[1] if len(record.args) > 1 else [],
                "final_answer": record.args[2] if len(record.args) > 2 else None
            }
            return json.dumps(log_record, ensure_ascii=False, indent=4)

    logger = logging.getLogger("text_logger")
    logger.setLevel(logging.INFO)
    handler = RotatingFileHandler(
        temp_dir_path / "{}_AP_enhanced_langchain.log".format(model_name)
    )
    handler.setFormatter(JsonFormatter())
    logger.addHandler(handler)

    return temp_dir_path, logger


def init_chat_logger():
    """Initialize chat logger for .chat file like AgentScope"""
    global temp_dir_path
    chat_log_path = temp_dir_path / "{}_AP_enhanced_langchain.chat".format(model_name)
    return chat_log_path


def save_chat_message(chat_log_path, message_data):
    """Save a single chat message to .chat file in AgentScope format"""
    import time
    from datetime import datetime
    import uuid
    
    # Format message in AgentScope style
    chat_record = {
        "__module__": "langchain.schema.messages",
        "__name__": "ChatMessage", 
        "id": str(uuid.uuid4()).replace('-', ''),
        "name": message_data.get('name', 'langchain_agent'),
        "role": message_data.get('role', 'assistant'),
        "content": message_data.get('content', []),
        "metadata": message_data.get('metadata', None),
        "timestamp": datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }
    
    # Append to chat file (one JSON per line, like AgentScope)
    with open(chat_log_path, 'a', encoding='utf-8') as f:
        f.write(json.dumps(chat_record, ensure_ascii=False) + '\n')


def load_langchain_config(config_path='./agent/config_Kimik2.json'):
    """Load configuration and initialize LangChain components"""
    with open(config_path, 'r') as f:
        config = json.load(f)
    
    # Initialize OpenAI model with stricter parameters
    model_config = config['models'][0]
    llm_kwargs = {
        'model': model_config['model_name'],
        'api_key': model_config['api_key'],
        'base_url': model_config['client_args']['base_url'],
        'temperature': 0.1,  # Lower temperature for more focused responses
        'request_timeout': 120  # 2 minute timeout per request
    }
    
    # Add generate_args via extra_body if present in config
    if 'generate_args' in model_config:
        llm_kwargs['extra_body'] = model_config['generate_args']
    
    llm = ChatOpenAI(**llm_kwargs)
    
    # Prepare MCP servers configuration
    mcp_servers = {}
    for server_name, server_config in config['mcpServers'].items():
        # Update paths to use current temp directory
        updated_args = []
        for arg in server_config['args']:
            if 'tmp/tmp/out' in arg:
                updated_args.append(str(temp_dir_path / 'out'))
            elif arg.startswith('tools/'):
                updated_args.append('agent/' + arg)
            else:
                updated_args.append(arg)
        
        mcp_servers[server_name] = {
            "command": server_config['command'],
            "args": updated_args,
            "transport": "stdio"
        }
    
    return llm, mcp_servers


async def create_langchain_agent(llm, mcp_servers):
    """Create LangChain ReAct agent with MCP tools"""
    # Create MCP client
    client = MultiServerMCPClient(mcp_servers)
    
    try:
        # Get tools from all MCP servers
        tools = await client.get_tools()
        print(f"Successfully loaded {len(tools)} tools from MCP servers")
        
        # Create ReAct agent
        agent = create_react_agent(llm, tools)
        
        return agent, client
    except Exception as e:
        print(f"Error creating agent: {e}")
        if hasattr(client, 'close'):
            await client.close()
        raise


def load_questions(test_json_path: str = 'benchmark/question.json'):
    """Load evaluation questions"""
    with open(test_json_path, 'r') as f:
        test_json = json.load(f)

    out = []
    for _, (question_idx, question_info) in enumerate(test_json.items()):
        AP_INDEX = 0 if question_info['evaluation'][0]['type'] == 'autonomous planning' else 1
        data = question_info['evaluation'][AP_INDEX].get('data', None)
        data = question_info['evaluation'][1 - AP_INDEX].get('data', None) if data is None else data

        if data is None:
            continue
        out.append({
            "question_id": question_idx,
            "auto": question_info['evaluation'][AP_INDEX]['question'],
            "instruct": question_info['evaluation'][1 - AP_INDEX]['question'],
            "data": data,
            "choices": question_info.get('choices', None)
        })

    return out


def extract_answer_from_response(response):
    """Extract final answer from agent response"""
    messages = response.get("messages", [])
    
    # Look for the final answer in the last assistant message
    for message in reversed(messages):
        if hasattr(message, 'type') and message.type == 'ai':
            content = message.content
            if '<Answer>' in content and '</Answer>' in content:
                # Extract answer between tags
                start = content.find('<Answer>') + len('<Answer>')
                end = content.find('</Answer>')
                return content[start:end].strip()
            return content
    
    return "No answer found"


async def handle_question(agent, question, chat_log_path=None):
    """Handle a single question with the LangChain agent"""
    try:
        # Prepare query
        query = question['auto'] + question['data'] if autoplanning else \
            question['instruct'] + question['data']

        if question['choices']:
            query += '\n'.join([''] + [
                '{}.{}'.format(chr(ord('A') + i), choice) 
                for i, choice in enumerate(question['choices'])
            ])

        # Add system prompt
        full_query = f"{sys_prompt}\n\nQuestion: {query}"
        
        print(f"\n--- Processing Question {question['question_id']} ---")
        print(f"Query: {query[:200]}...")

        # Invoke agent with configuration to prevent infinite loops
        response = await agent.ainvoke(
            {"messages": [HumanMessage(content=full_query)]},
            config={
                "recursion_limit": 50,  # Increase recursion limit
                "max_execution_time": 300,  # 5 minutes timeout
            }
        )
        
        # Extract final answer
        final_answer = extract_answer_from_response(response)
        
        # Convert response to AgentScope-compatible format for logging
        conversation_log = []
        
        for message in response.get("messages", []):
            if hasattr(message, 'type'):
                if message.type == 'human':
                    # User message
                    conversation_log.append({
                        "role": "user",
                        "content": message.content
                    })
                
                elif message.type == 'ai':
                    # Assistant message - handle both thinking and tool calls
                    assistant_content = []
                    
                    # First check if there's thinking content (text before tool calls)
                    if message.content and message.content.strip():
                        assistant_content.append({
                            "type": "text",
                            "content": message.content
                        })
                    
                    # Then check for tool calls
                    if hasattr(message, 'additional_kwargs') and 'tool_calls' in message.additional_kwargs:
                        # Format tool calls in AgentScope style
                        for tool_call in message.additional_kwargs['tool_calls']:
                            try:
                                arguments = json.loads(tool_call['function']['arguments']) if isinstance(tool_call['function']['arguments'], str) else tool_call['function']['arguments']
                            except:
                                arguments = tool_call['function']['arguments']
                            
                            assistant_content.append({
                                "name": tool_call['function']['name'],
                                "input": arguments
                            })
                    
                    # Only add to log if there's actual content
                    if assistant_content:
                        conversation_log.append({
                            "role": "assistant", 
                            "content": assistant_content
                        })
                
                elif message.type == 'tool':
                    # Tool result message in AgentScope format
                    conversation_log.append({
                        "role": "tool",
                        "name": message.name,
                        "content": [{
                            "output": [{
                                "text": str(message.content),
                            }],
                        }]
                    })
        
        
        # Log the conversation in the same format as original code (only to .log file)
        logger.info("Chat Content", question['question_id'], conversation_log, final_answer)
        
        print(f"Final Answer: {final_answer}")
        return final_answer
        
    except Exception as e:
        error_msg = f"Error processing question {question['question_id']}: {e}"
        print(error_msg)

        logger.info(question['question_id'], [], error_msg)
        return f"Error: {e}"


async def main():
    """Main evaluation function"""
    # Parse command line arguments
    args = parse_args()

    # Load enhanced config (required)
    load_enhanced_config(args.enhanced_config)

    print(f"\n{'='*80}")
    print("Initializing Enhanced Earth Science Agent")
    print(f"{'='*80}\n")

    # Initialize global parameters
    init_global_params(batch_dir=args.batch_dir)

    # Initialize chat logger - removed, only keep .log files
    # chat_log_path = init_chat_logger()

    print(f"Output directory: {temp_dir_path}")
    if args.batch_dir:
        print(f"Batch mode: Processing specific question subset")
    print(f"Mode: Enhanced with {num_experiences} learned experiences\n")

    # Load configuration and create agent
    llm, mcp_servers = load_langchain_config()
    agent, client = await create_langchain_agent(llm, mcp_servers)
    
    try:
        # Load questions
        all_questions = load_questions()

        # Evaluation question IDs (188 questions - excluding 60 training questions)
        target_question_ids = [
            '1', '5', '6', '7', '9', '13', '14', '15', '16', '18', '20', '22', '24', '25', '26',
            '27', '29', '31', '32', '34', '35', '36', '37', '38', '39', '41', '42', '44', '46',
            '47', '48', '49', '50', '51', '52', '54', '56', '57', '58', '59', '60', '61', '62',
            '63', '64', '65', '66', '67', '68', '69', '70', '71', '72', '73', '74', '75', '76',
            '77', '78', '79', '81', '82', '83', '84', '85', '86', '87', '88', '89', '90', '91',
            '92', '93', '94', '95', '96', '97', '98', '99', '100', '101', '103', '104', '105',
            '106', '108', '111', '113', '115', '117', '118', '120', '122', '123', '124', '126',
            '127', '128', '129', '130', '131', '132', '133', '134', '135', '139', '140', '141',
            '142', '143', '144', '145', '146', '147', '148', '150', '151', '152', '153', '154',
            '156', '157', '158', '159', '160', '162', '163', '164', '165', '166', '167', '168',
            '169', '170', '171', '172', '174', '175', '177', '178', '179', '180', '181', '182',
            '183', '184', '185', '188', '189', '192', '193', '194', '196', '197', '198', '199',
            '200', '201', '204', '205', '206', '207', '211', '212', '213', '214', '216', '217',
            '219', '220', '221', '222', '226', '227', '230', '232', '234', '235', '237', '238',
            '239', '240', '241', '243', '245', '246', '247', '248'
        ]

        # If specific question IDs are provided, use them; otherwise use all
        if args.question_ids:
            target_question_ids = args.question_ids
            print(f"Batch mode: Evaluating {len(target_question_ids)} specific questions")

        # Filter questions based on target IDs
        questions = [q for q in all_questions if q['question_id'] in target_question_ids]

        print(f"Loaded {len(questions)} questions for evaluation (out of {len(all_questions)} total)")
        
        # Process questions
        results = []
        for question in tqdm(questions, desc="Processing questions"):
            # Removed chat_log_path parameter
            answer = await handle_question(agent, question, None)
            results.append({
                "question_id": question['question_id'],
                "answer": answer
            })
            
            # Optional: Add delay between questions to avoid rate limiting
            await asyncio.sleep(1)
        
        # Save results summary
        results_path = temp_dir_path / "results_summary.json"
        with open(results_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=4)
        
        print(f"\nEvaluation completed! Results saved to {results_path}")
        print(f"Detailed logs available at: {temp_dir_path}")
        
    except Exception as e:
        print(f"Error in main evaluation: {e}")
        raise
    
    finally:
        # Clean up
        if hasattr(client, 'close'):
            await client.close()


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