from fastmcp import FastMCP
import os

# 创建 MCP 服务
mcp = FastMCP("Zspace_Knowledge")

@mcp.tool()
def search_local_knowledge(query: str) -> str:
    """
    当用户询问药房、家庭物资分布等信息时，调用此工具读取极空间本地知识库寻找答案。
    """
    data_dir = "/data"
    knowledge_text = ""
    
    if not os.path.exists(data_dir):
        return "系统提示：未找到极空间的知识库文件夹。"
        
    # 读取文件夹下所有的 txt 文档
    for filename in os.listdir(data_dir):
        if filename.endswith(".txt"):
            filepath = os.path.join(data_dir, filename)
            try:
                with open(filepath, 'r', encoding='utf-8') as f:
                    knowledge_text += f"\n--- 【{filename}】 ---\n"
                    knowledge_text += f.read()
            except Exception as e:
                pass
                
    if not knowledge_text:
        return "系统提示：知识库为空，请在极空间文件夹中放入 txt 文档。"
        
    return f"请根据以下极空间私有知识库内容，回答用户的问题 '{query}'：\n\n{knowledge_text}"

if __name__ == "__main__":
    # 使用标准输出模式与桥接脚本通信
    mcp.run(transport='stdio')