TravelContentCreator/examples/test_stream.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import json
import time
import logging

# Determine the project root directory (assuming examples/ is one level down)
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.append(PROJECT_ROOT)

# Now import from core
try:
    from core.ai_agent import AI_Agent
except ImportError as e:
    logging.critical(f"Failed to import AI_Agent. Ensure '{PROJECT_ROOT}' is in sys.path and core/ai_agent.py exists. Error: {e}")
    sys.exit(1)

def load_config(config_path):
    """Loads configuration from a JSON file."""
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        logging.info(f"Config loaded successfully from {config_path}")
        return config
    except FileNotFoundError:
        logging.error(f"Error: Configuration file not found at {config_path}")
        return None
    except json.JSONDecodeError:
        logging.error(f"Error: Could not decode JSON from {config_path}")
        return None
    except Exception as e:
        logging.exception(f"An unexpected error occurred loading config {config_path}:")
        return None

def main():
    # --- Basic Logging Setup --- 
    logging.basicConfig(
        level=logging.INFO, 
        format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    # --- End Logging Setup --- 

    logging.info("Starting AI Agent Stream Test...")

    # Load configuration (adjust path relative to this script)
    config_path = os.path.join(PROJECT_ROOT, "poster_gen_config.json")
    config = load_config(config_path)
    if config is None:
        logging.critical("Failed to load configuration. Exiting test.")
        sys.exit(1)

    # Example Prompts
    system_prompt = "你是一个乐于助人的AI助手，擅长写短篇故事。"
    user_prompt = "请写一个关于旅行机器人的短篇故事，它在一个充满异国情调的星球上发现了新的生命形式。"

    ai_agent = None
    try:
        # --- Extract AI Agent parameters from config --- 
        ai_api_url = config.get("api_url")
        ai_model = config.get("model")
        ai_api_key = config.get("api_key")
        request_timeout = config.get("request_timeout", 30)
        max_retries = config.get("max_retries", 3)
        stream_chunk_timeout = config.get("stream_chunk_timeout", 60) # Get stream chunk timeout
        
        # Check for required AI params
        if not all([ai_api_url, ai_model, ai_api_key]):
             logging.critical("Missing required AI configuration (api_url, model, api_key) in config. Exiting test.")
             sys.exit(1)
        # --- End Extract AI Agent params --- 

        logging.info("Initializing AI Agent for stream test...")
        # Initialize AI_Agent using extracted parameters
        ai_agent = AI_Agent(
            api_url=ai_api_url,     # Use extracted var
            model=ai_model,      # Use extracted var
            api_key=ai_api_key,       # Use extracted var
            timeout=request_timeout,
            max_retries=max_retries,
            stream_chunk_timeout=stream_chunk_timeout # Pass it here
        )

        # Example call to work_stream
        logging.info("Calling ai_agent.work_stream...")
        # Extract generation parameters from config
        temperature = config.get("content_temperature", 0.7) # Use a relevant temperature setting
        top_p = config.get("content_top_p", 0.9)
        presence_penalty = config.get("content_presence_penalty", 0.0)

        start_time = time.time()
        stream_generator = ai_agent.work_stream(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            info_directory=None, # No extra context folder for this test
            temperature=temperature,
            top_p=top_p,
            presence_penalty=presence_penalty
        )

        # Process the stream
        logging.info("Processing stream response:")
        full_response = ""
        for chunk in stream_generator:
            print(chunk, end="", flush=True) # Keep print for stream output
            full_response += chunk

        end_time = time.time()
        logging.info(f"\n--- Stream Finished ---")
        logging.info(f"Total time: {end_time - start_time:.2f} seconds")
        logging.info(f"Total characters received: {len(full_response)}")

    except KeyError as e:
        logging.error(f"Configuration error: Missing key '{e}'. Please check '{config_path}'.")
    except Exception as e:
        logging.exception("An error occurred during the stream test:")
    finally:
        # Ensure the agent is closed
        if ai_agent:
            logging.info("Closing AI Agent...")
            ai_agent.close()
            logging.info("AI Agent closed.")

if __name__ == "__main__":
    main()