ai_architecture/examples/basic_usage.py

"""
基本使用示例
"""

import os
import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.core.llm_manager import llm_manager
from src.core.vector_store import vector_store_manager
from src.core.document_processor import document_processor
from src.services.chat_service import chat_service
from src.services.qa_service import qa_service
from src.services.document_service import document_service
from src.utils.logger import setup_logger
from loguru import logger


def test_llm_manager():
    """测试LLM管理器"""
    print("=== 测试LLM管理器 ===")
    
    # 列出可用的提供商
    providers = llm_manager.list_providers()
    print(f"可用的LLM提供商: {providers}")
    
    # 获取OpenRouter可用模型
    openrouter_models = llm_manager.get_available_models("openrouter")
    print(f"OpenRouter可用模型: {openrouter_models[:5]}...")  # 只显示前5个
    
    # 测试连接
    if llm_manager.test_connection("openrouter"):
        print("OpenRouter连接测试成功")
    else:
        print("OpenRouter连接测试失败")


def test_chat_service():
    """测试聊天服务"""
    print("\n=== 测试聊天服务 ===")
    
    # 创建会话
    session = chat_service.create_session("test_session", "你是一个有用的AI助手")
    print(f"创建会话: {session.session_id}")
    
    # 发送消息
    try:
        response = chat_service.send_message("test_session", "你好，请介绍一下自己")
        print(f"AI回复: {response['assistant_message']}")
    except Exception as e:
        print(f"聊天失败: {e}")
    
    # 列出会话
    sessions = chat_service.list_sessions()
    print(f"当前会话: {sessions}")


def test_document_service():
    """测试文档服务"""
    print("\n=== 测试文档服务 ===")
    
    # 创建测试文档
    test_doc_path = "test_document.txt"
    with open(test_doc_path, "w", encoding="utf-8") as f:
        f.write("这是一个测试文档。\n它包含一些示例文本。\n用于测试文档处理功能。")
    
    try:
        # 验证文档
        validation = document_service.validate_document(test_doc_path)
        print(f"文档验证结果: {validation}")
        
        if validation["valid"]:
            # 处理并存储文档
            result = document_service.process_and_store_document(
                test_doc_path, 
                collection_name="test_collection"
            )
            print(f"文档处理结果: {result}")
            
            # 搜索文档
            search_results = document_service.search_documents(
                "测试", 
                collection_name="test_collection"
            )
            print(f"搜索结果: {len(search_results)} 个文档")
            
    except Exception as e:
        print(f"文档处理失败: {e}")
    
    # 清理测试文件
    if os.path.exists(test_doc_path):
        os.remove(test_doc_path)


def test_qa_service():
    """测试问答服务"""
    print("\n=== 测试问答服务 ===")
    
    try:
        # 创建测试文档并添加到问答系统
        test_doc_path = "qa_test_doc.txt"
        with open(test_doc_path, "w", encoding="utf-8") as f:
            f.write("人工智能（AI）是计算机科学的一个分支，旨在创建能够执行通常需要人类智能的任务的系统。")
        
        # 处理文档
        documents = document_processor.process_document_pipeline(test_doc_path)
        qa_service.add_documents_for_qa(documents, collection_name="qa_test")
        
        # 提问
        question = "什么是人工智能？"
        answer = qa_service.ask_question(question, collection_name="qa_test")
        print(f"问题: {question}")
        print(f"答案: {answer['answer']}")
        print(f"来源数量: {answer['source_count']}")
        
        # 清理测试文件
        if os.path.exists(test_doc_path):
            os.remove(test_doc_path)
            
    except Exception as e:
        print(f"问答测试失败: {e}")


def main():
    """主函数"""
    print("AI Architecture 基本使用示例")
    print("=" * 50)
    
    # 设置日志
    setup_logger()
    
    # 检查环境变量
    if not os.getenv("OPENROUTER_API_KEY"):
        print("警告: 未设置OPENROUTER_API_KEY环境变量")
        print("请复制env.example为.env并设置您的OpenRouter API密钥")
        return
    
    # 运行测试
    test_llm_manager()
    test_chat_service()
    test_document_service()
    test_qa_service()
    
    print("\n示例运行完成！")


if __name__ == "__main__":
    main()